Improve Selector usability

This commit improves the usability of Selectors by adding new methods to
the Selector interface and by deprecating its runner functionality. The
new methods include methods to stream matching shapes, stream matches,
and consume matches. Streaming matches and shapes requires a minimal
amount of buffering, and they provide the ability to more easily perform
checks like determine if any shapes in a model match a selector. If all
matching shapes are necessary, calling select is preferred. If consuming
matches can be done in a way that is push-based (e.g., writing to CLI
output), then using consumeMatches is preferred to matches.

The previous runner() method of Selector is now deprecated and just
delegates calls back to the Selector. runner() was initially intended to
allow for the customization of the environment in which a selector is
executed, but this design turned out to be unnecessary, and as such, the
need to create a runner and use a kind of builder interface to use
Selectors made them awkward to work with for no benefit.

A final change is made to implement an IdentitySelector to optimize
matching all shapes in a model.

smithy-cli/src/main/java/software/amazon/smithy/cli/commands/SelectCommand.java

@@ -87,10 +87,10 @@ public void execute(Arguments arguments, ClassLoader classLoader) {
         } else {
             // Show the JSON output for writing with --vars.
             List<Node> result = new ArrayList<>();
-            selector.runner().model(model).selectMatches((shape, vars) -> {
+            selector.consumeMatches(model, match -> {
                 result.add(Node.objectNodeBuilder()
-                        .withMember("shape", Node.from(shape.getId().toString()))
-                        .withMember("vars", collectVars(vars))
+                        .withMember("shape", Node.from(match.getShape().getId().toString()))
+                        .withMember("vars", collectVars(match))
                         .build());
             });
             Cli.stdout(Node.prettyPrintJson(new ArrayNode(result, SourceLocation.NONE)));

smithy-model/src/main/java/software/amazon/smithy/model/loader/ValidatorDefinition.java

@@ -54,10 +54,8 @@ List<ValidationEvent> map(Model model, List<ValidationEvent> events) {
 
         // If there's a selector, create a list of candidate shape IDs that can be emitted.
         if (selector != null) {
-            candidates = selector.runner()
-                    .model(model)
-                    .selectShapes()
-                    .stream()
+            candidates = selector
+                    .shapes(model)
                     .map(Shape::getId)
                     .collect(Collectors.toSet());
         }

smithy-model/src/main/java/software/amazon/smithy/model/selector/Context.java

@@ -20,15 +20,14 @@
 import java.util.Set;
 import software.amazon.smithy.model.knowledge.NeighborProviderIndex;
 import software.amazon.smithy.model.shapes.Shape;
-import software.amazon.smithy.utils.MapUtils;
 
 /**
  * Selector evaluation context object.
  */
 final class Context {
 
     NeighborProviderIndex neighborIndex;
-    private Map<String, Set<Shape>> variables;
+    private final Map<String, Set<Shape>> variables;
 
     Context(NeighborProviderIndex neighborIndex) {
         this.neighborIndex = neighborIndex;
@@ -45,15 +44,6 @@ Context clearVars() {
         return this;
     }
 
-    /**
-     * Copies the current set of variables to an immutable map.
-     *
-     * @return Returns a copy of the variables.
-     */
-    Map<String, Set<Shape>> copyVars() {
-        return MapUtils.copyOf(variables);
-    }
-
     /**
      * Gets the currently set variables.
      *

smithy-model/src/main/java/software/amazon/smithy/model/selector/IdentitySelector.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.model.selector;
+
+import java.util.Collections;
+import java.util.Set;
+import java.util.stream.Stream;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.Shape;
+
+/**
+ * An optimized Selector implementation that uses the provided Model directly
+ * rather than needing to send each shape through the Selector machinery.
+ *
+ * @see Selector#IDENTITY
+ */
+final class IdentitySelector implements Selector {
+    @Override
+    public Set<Shape> select(Model model) {
+        return model.toSet();
+    }
+
+    @Override
+    public Stream<Shape> shapes(Model model) {
+        return model.shapes();
+    }
+
+    @Override
+    public Stream<ShapeMatch> matches(Model model) {
+        return model.shapes().map(shape -> new ShapeMatch(shape, Collections.emptyMap()));
+    }
+
+    @Override
+    public String toString() {
+        return "*";
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        return other instanceof Selector && toString().equals(other.toString());
+    }
+
+    @Override
+    public int hashCode() {
+        return toString().hashCode();
+    }
+}

smithy-model/src/main/java/software/amazon/smithy/model/selector/Selector.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License").
  * You may not use this file except in compliance with the License.
@@ -15,25 +15,26 @@
 
 package software.amazon.smithy.model.selector;
 
-import java.util.HashSet;
+import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.SourceException;
-import software.amazon.smithy.model.knowledge.NeighborProviderIndex;
 import software.amazon.smithy.model.node.Node;
 import software.amazon.smithy.model.shapes.Shape;
-import software.amazon.smithy.utils.ListUtils;
-import software.amazon.smithy.utils.SmithyBuilder;
 
 /**
  * Matches a set of shapes using a selector expression.
  */
 public interface Selector {
 
     /** A selector that always returns all provided values. */
-    Selector IDENTITY = new WrappedSelector("*", ListUtils.of(InternalSelector.IDENTITY));
+    Selector IDENTITY = new IdentitySelector();
 
     /**
      * Parses a selector expression.
@@ -42,7 +43,11 @@ public interface Selector {
      * @return Returns the parsed {@link Selector}.
      */
     static Selector parse(String expression) {
-        return SelectorParser.parse(expression);
+        if (expression.equals("*")) {
+            return IDENTITY;
+        } else {
+            return SelectorParser.parse(expression);
+        }
     }
 
     /**
@@ -67,7 +72,65 @@ static Selector fromNode(Node node) {
      * @return Returns the matching shapes.
      */
     default Set<Shape> select(Model model) {
-        return runner().model(model).selectShapes();
+        return shapes(model).collect(Collectors.toSet());
+    }
+
+    /**
+     * Matches a selector to a set of shapes and receives each matched shape
+     * with the variables that were set when the shape was matched.
+     *
+     * @param model Model to select shapes from.
+     * @param shapeMatchConsumer Receives each matched shape and the vars available when the shape was matched.
+     */
+    default void consumeMatches(Model model, Consumer<ShapeMatch> shapeMatchConsumer) {
+        matches(model).forEach(shapeMatchConsumer);
+    }
+
+    /**
+     * Returns a stream of shapes in a model that match the selector.
+     *
+     * @param model Model to match the selector against.
+     * @return Returns a stream of matching shapes.
+     */
+    Stream<Shape> shapes(Model model);
+
+    /**
+     * Returns a stream of {@link ShapeMatch} objects for each match found in
+     * a model.
+     *
+     * @param model Model to match the selector against.
+     * @return Returns a stream of {@code ShapeMatch} objects.
+     */
+    Stream<ShapeMatch> matches(Model model);
+
+    /**
+     * Represents a selector match found in the model.
+     *
+     * <p>The {@code getShape} method is used to get the shape that matched,
+     * and all of the contextual variables that were set when the match
+     * occurred can be accessed using typical {@link Map} methods like
+     * {@code get}, {@code contains}, etc.
+     */
+    final class ShapeMatch extends HashMap<String, Set<Shape>> {
+        private final Shape shape;
+
+        /**
+         * @param shape Shape that matched.
+         * @param variables Variables that matched. This map is copied into ShapeMatch.
+         */
+        public ShapeMatch(Shape shape, Map<String, Set<Shape>> variables) {
+            super(variables);
+            this.shape = shape;
+        }
+
+        /**
+         * Gets the matching shape.
+         *
+         * @return Returns the matching shape.
+         */
+        public Shape getShape() {
+            return shape;
+        }
     }
 
     /**
@@ -76,79 +139,43 @@ default Set<Shape> select(Model model) {
      *
      * @return Returns the created runner.
      */
-    Runner runner();
+    @Deprecated
+    default Runner runner() {
+        return new Runner(this);
+    }
 
     /**
      * Builds the execution environment for a selector and executes selectors.
+     *
+     * @deprecated This class is no longer necessary. It was originally intended
+     *  to allow more customization to how selectors are executed against a model,
+     *  but this has proven unnecessary.
      */
+    @Deprecated
     final class Runner {
 
-        private final InternalSelector delegate;
-        private final Class<? extends Shape> startingShapeType;
+        private final Selector selector;
         private Model model;
 
-        Runner(InternalSelector delegate, Class<? extends Shape> startingShapeType) {
-            this.delegate = delegate;
-            this.startingShapeType = startingShapeType;
+        Runner(Selector selector) {
+            this.selector = selector;
         }
 
-        /**
-         * Sets the <em>required</em> model to use to select shapes with.
-         *
-         * @param model Model used in the selector evaluation.
-         * @return Returns the Runner.
-         */
+        @Deprecated
         public Runner model(Model model) {
             this.model = model;
             return this;
         }
 
-        /**
-         * Runs the selector and returns the set of matching shapes.
-         *
-         * @return Returns the set of matching shapes.
-         * @throws IllegalStateException if a {@code model} has not been set.
-         */
+        @Deprecated
         public Set<Shape> selectShapes() {
-            Set<Shape> result = new HashSet<>();
-            pushShapes((ctx, s) -> {
-                result.add(s);
-                return true;
-            });
-            return result;
-        }
-
-        private Context createContext() {
-            SmithyBuilder.requiredState("model", model);
-            return new Context(NeighborProviderIndex.of(model));
+            return selector.select(Objects.requireNonNull(model, "model not set"));
         }
 
-        /**
-         * Matches a selector to a set of shapes and receives each matched shape
-         * with the variables that were set when the shape was matched.
-         *
-         * @param matchConsumer Receives each matched shape and the vars available when the shape was matched.
-         * @throws IllegalStateException if a {@code model} has not been set.
-         */
+        @Deprecated
         public void selectMatches(BiConsumer<Shape, Map<String, Set<Shape>>> matchConsumer) {
-            pushShapes((ctx, s) -> {
-                matchConsumer.accept(s, ctx.copyVars());
-                return true;
-            });
-        }
-
-        private void pushShapes(InternalSelector.Receiver acceptor) {
-            Context context = createContext();
-
-            if (startingShapeType != null) {
-                model.shapes(startingShapeType).forEach(shape -> {
-                    delegate.push(context.clearVars(), shape, acceptor);
-                });
-            } else {
-                for (Shape shape : model.toSet()) {
-                    delegate.push(context.clearVars(), shape, acceptor);
-                }
-            }
+            selector.consumeMatches(Objects.requireNonNull(model, "model not set"),
+                                    m -> matchConsumer.accept(m.getShape(), m));
         }
     }
 }