Add "trait" relationship to selectors

A "trait" relationship allows selectors to traverse from shapes to the
the trait shapes applied to the shape. This can be used to query for
things like "services with no protocol traits", "deprecated traits that
are being used on shapes", "traits that aren't referenced by any
shapes", etc.

Trait relationships are not typically needed and add a lot of overhead.
As such, they are opt-in only and enabled in selectors by adding a named
"trait" directed relationship. They're also opt-in in code.

In order to pull this off, I updated the selector API to always require
a Model.

Author: mtdowling

docs/source/spec/core/selectors.rst

@@ -258,6 +258,17 @@ an identifier:
 
     resource:test(-[identifier]->)
 
+Relationships from a shape to the traits applied to the shape can be traversed
+using a directed relationship named ``trait``. It is atypical to traverse
+``trait`` relationships, therefore they are only yielded by selectors when
+explicitly requested using a ``trait`` directed relationship. The following
+selector finds all service shapes that have a protocol trait applied to it
+(that is, a trait that is marked with the :ref:`protocolDefinition-trait`):
+
+::
+
+    service:test(-[trait]-> [trait|protocolDefinition])
+
 
 .. _selector-relationships:
 
@@ -353,6 +364,12 @@ The table below lists the labeled directed relationships from each shape.
       -
       - The shape targeted by the member. Note that member targets have no
         relationship name.
+    * - ``*``
+      - trait
+      - Each trait applied to a shape. The neighbor shape is the shape that
+        defines the trait. This kind of relationship is only traversed if the
+        ``trait`` relationship is explicitly stated as a desired directed
+        neighbor relationship type.
 
 .. important::
 
@@ -496,6 +513,20 @@ the member does not have the ``length`` trait:
         :test(> string:not([trait|length]))
         :test(:not([trait|length]))
 
+The following selector finds all service shapes that do not have a
+protocol trait applied to it:
+
+::
+
+    service:not(:test(-[trait]-> [trait|protocolDefinition]))
+
+The following selector finds all traits that are not attached to any shape
+in the model:
+
+::
+
+    :not(* -[trait]-> *)[trait|trait]
+
 
 :of
 ~~~
@@ -576,6 +607,7 @@ Selectors are defined by the following ABNF_ grammar.
                          :/ "instanceOperation"
                          :/ "resource"
                          :/ "bound"
+                         :/ "trait"
     attr                   :"[" `attr_key` *(`comparator` `attr_value` ["i"]) "]"
     attr_key               :`id_attribute` / `trait_attribute` / `service_attribute`
     id_attribute           :"id" ["|" ("namespace" / "name" / "member")]

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

@@ -57,7 +57,7 @@ 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) {
             NeighborProvider provider = model.getKnowledge(NeighborProviderIndex.class).getProvider();
-            candidates = selector.select(provider, model).stream()
+            candidates = selector.select(model, provider).stream()
                     .map(Shape::getId)
                     .collect(Collectors.toSet());
         }

smithy-model/src/main/java/software/amazon/smithy/model/neighbor/NeighborProvider.java

@@ -15,6 +15,7 @@
 
 package software.amazon.smithy.model.neighbor;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -57,6 +58,36 @@ static NeighborProvider precomputed(Model model) {
         return precomputed(model, of(model));
     }
 
+    /**
+     * Creates a NeighborProvider that includes {@link RelationshipType#TRAIT}
+     * relationships.
+     *
+     * @param model Model to use to look up trait shapes.
+     * @param neighborProvider Provider to wrap.
+     * @return Returns the created neighbor provider.
+     */
+    static NeighborProvider withTraitRelationships(Model model, NeighborProvider neighborProvider) {
+        return shape -> {
+            List<Relationship> relationships = neighborProvider.getNeighbors(shape);
+
+            // Don't copy the array unless the shape has traits.
+            if (shape.getAllTraits().isEmpty()) {
+                return relationships;
+            }
+
+            // The delegate might have returned an immutable list, so copy first.
+            relationships = new ArrayList<>(relationships);
+            for (ShapeId trait : shape.getAllTraits().keySet()) {
+                Relationship traitRel = model.getShape(trait)
+                        .map(target -> Relationship.create(shape, RelationshipType.TRAIT, target))
+                        .orElseGet(() -> Relationship.createInvalid(shape, RelationshipType.TRAIT, trait));
+                relationships.add(traitRel);
+            }
+
+            return relationships;
+        };
+    }
+
     /**
      * Creates a NeighborProvider that precomputes the neighbors of a model.
      *

smithy-model/src/main/java/software/amazon/smithy/model/neighbor/RelationshipType.java

@@ -16,6 +16,7 @@
 package software.amazon.smithy.model.neighbor;
 
 import java.util.Optional;
+import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.selector.Selector;
 import software.amazon.smithy.model.shapes.ListShape;
 import software.amazon.smithy.model.shapes.MapShape;
@@ -25,6 +26,7 @@
 import software.amazon.smithy.model.shapes.SetShape;
 import software.amazon.smithy.model.shapes.StructureShape;
 import software.amazon.smithy.model.shapes.UnionShape;
+import software.amazon.smithy.model.traits.TraitDefinition;
 
 /**
  * Defines the relationship types between neighboring shapes.
@@ -187,7 +189,19 @@ public enum RelationshipType {
      * shapes. They reference the {@link MemberShape member} shapes that define
      * the members of the union.
      */
-    UNION_MEMBER("member", RelationshipDirection.DIRECTED);
+    UNION_MEMBER("member", RelationshipDirection.DIRECTED),
+
+    /**
+     * Relationships that exist between a shape and traits bound to the
+     * shape. They reference shapes marked with the {@link TraitDefinition}
+     * trait.
+     *
+     * <p>This kind of relationship is not returned by default from a
+     * {@link NeighborProvider}. You must explicitly wrap a {@link NeighborProvider}
+     * with {@link NeighborProvider#withTraitRelationships(Model, NeighborProvider)}
+     * in order to yield trait relationships.
+     */
+    TRAIT("trait", RelationshipDirection.DIRECTED);
 
     private String selectorLabel;
     private RelationshipDirection direction;

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

@@ -17,6 +17,7 @@
 
 import java.util.List;
 import java.util.Set;
+import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.neighbor.NeighborProvider;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.utils.SetUtils;
@@ -40,9 +41,9 @@ static Selector of(List<Selector> predicates) {
     }
 
     @Override
-    public Set<Shape> select(NeighborProvider neighborProvider, Set<Shape> shapes) {
+    public Set<Shape> select(Model model, NeighborProvider neighborProvider, Set<Shape> shapes) {
         for (Selector selector : selectors) {
-            shapes = selector.select(neighborProvider, shapes);
+            shapes = selector.select(model, neighborProvider, shapes);
             if (shapes.isEmpty()) {
                 return SetUtils.of();
             }

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

@@ -22,6 +22,7 @@
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.neighbor.NeighborProvider;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.Shape;
@@ -75,7 +76,7 @@ interface Comparator extends BiFunction<String, String, Boolean> {}
     }
 
     @Override
-    public Set<Shape> select(NeighborProvider neighborProvider, Set<Shape> shapes) {
+    public Set<Shape> select(Model model, NeighborProvider neighborProvider, Set<Shape> shapes) {
         return shapes.stream()
                 .filter(shape -> matchesAttribute(key.apply(shape)))
                 .collect(Collectors.toSet());

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

@@ -18,6 +18,7 @@
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
+import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.neighbor.NeighborProvider;
 import software.amazon.smithy.model.shapes.Shape;
 
@@ -36,9 +37,9 @@ static Selector of(List<Selector> predicates) {
     }
 
     @Override
-    public Set<Shape> select(NeighborProvider neighborProvider, Set<Shape> shapes) {
+    public Set<Shape> select(Model model, NeighborProvider neighborProvider, Set<Shape> shapes) {
         return selectors.stream()
-                .flatMap(selector -> selector.select(neighborProvider, shapes).stream())
+                .flatMap(selector -> selector.select(model, neighborProvider, shapes).stream())
                 .collect(Collectors.toSet());
     }
 }

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

@@ -15,47 +15,62 @@
 
 package software.amazon.smithy.model.selector;
 
+import java.util.HashSet;
 import java.util.List;
-import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
+import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.neighbor.NeighborProvider;
 import software.amazon.smithy.model.neighbor.Relationship;
 import software.amazon.smithy.model.neighbor.RelationshipType;
 import software.amazon.smithy.model.shapes.Shape;
-import software.amazon.smithy.utils.OptionalUtils;
 
 /**
  * Traverses into the neighbors of shapes with an optional list of
  * neighbor rel filters.
  */
 final class NeighborSelector implements Selector {
+
     private final List<String> relTypes;
+    private final boolean includeTraits;
 
     NeighborSelector(List<String> relTypes) {
         this.relTypes = relTypes;
+        includeTraits = relTypes.contains("trait");
     }
 
     @Override
-    public Set<Shape> select(NeighborProvider neighborProvider, Set<Shape> shapes) {
-        return shapes.stream()
-                .flatMap(shape -> neighborProvider.getNeighbors(shape).stream().flatMap(this::mapNeighbor))
-                .collect(Collectors.toSet());
+    public Set<Shape> select(Model model, NeighborProvider neighborProvider, Set<Shape> shapes) {
+        NeighborProvider resolvedProvider = createProvider(model, neighborProvider);
+
+        Set<Shape> result = new HashSet<>();
+        for (Shape shape : shapes) {
+            for (Relationship rel : resolvedProvider.getNeighbors(shape)) {
+                if (rel.getNeighborShape().isPresent()) {
+                    Shape target = createNeighbor(rel, rel.getNeighborShape().get());
+                    if (target != null) {
+                        result.add(target);
+                    }
+                }
+            }
+        }
+
+        return result;
     }
 
-    private Stream<Shape> mapNeighbor(Relationship rel) {
-        return OptionalUtils.stream(rel.getNeighborShape()
-                .flatMap(target -> createNeighbor(rel, target)));
+    // Enable trait relationships only if explicitly asked for in a selector.
+    private NeighborProvider createProvider(Model model, NeighborProvider neighborProvider) {
+        return includeTraits
+               ? NeighborProvider.withTraitRelationships(model, neighborProvider)
+               : neighborProvider;
     }
 
-    private Optional<Shape> createNeighbor(Relationship rel, Shape target) {
+    private Shape createNeighbor(Relationship rel, Shape target) {
         if (rel.getRelationshipType() != RelationshipType.MEMBER_CONTAINER
                 && (relTypes.isEmpty() || relTypes.contains(getRelType(rel)))) {
-            return Optional.of(target);
+            return target;
         }
 
-        return Optional.empty();
+        return null;
     }
 
     private static String getRelType(Relationship rel) {

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

@@ -18,6 +18,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.neighbor.NeighborProvider;
 import software.amazon.smithy.model.shapes.Shape;
 
@@ -32,10 +33,10 @@ final class NotSelector implements Selector {
     }
 
     @Override
-    public Set<Shape> select(NeighborProvider neighborProvider, Set<Shape> shapes) {
+    public Set<Shape> select(Model model, NeighborProvider neighborProvider, Set<Shape> shapes) {
         Set<Shape> result = new HashSet<>(shapes);
         for (Selector predicate : selectors) {
-            result.removeAll(predicate.select(neighborProvider, shapes));
+            result.removeAll(predicate.select(model, neighborProvider, shapes));
         }
         return result;
     }

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

@@ -19,6 +19,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
+import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.neighbor.NeighborProvider;
 import software.amazon.smithy.model.neighbor.RelationshipType;
 import software.amazon.smithy.model.shapes.Shape;
@@ -44,7 +45,7 @@ final class OfSelector implements Selector {
     }
 
     @Override
-    public Set<Shape> select(NeighborProvider neighborProvider, Set<Shape> shapes) {
+    public Set<Shape> select(Model model, NeighborProvider neighborProvider, Set<Shape> shapes) {
         Set<Shape> result = new HashSet<>();
 
         // Filter out non-member shapes, and member shapes that cannot
@@ -55,7 +56,7 @@ public Set<Shape> select(NeighborProvider neighborProvider, Set<Shape> shapes) {
                 // If the parent provides a result for the predicate, then the
                 // Shape is not filtered out.
                 boolean anyMatch = selectors.stream()
-                        .anyMatch(selector -> !selector.select(neighborProvider, parentSet).isEmpty());
+                        .anyMatch(selector -> !selector.select(model, neighborProvider, parentSet).isEmpty());
                 if (anyMatch) {
                     result.add(shape);
                 }

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

@@ -110,7 +110,7 @@ public List<Path> search(ToShapeId startingShape, Selector targetSelector) {
         }
 
         // Find all shapes that match the selector then work backwards from there.
-        Set<Shape> candidates = targetSelector.select(neighborProvider, model);
+        Set<Shape> candidates = targetSelector.select(model, neighborProvider);
         if (candidates.isEmpty()) {
             LOGGER.info(() -> "No shapes matched the PathFinder selector of `" + targetSelector + "`");
             return ListUtils.of();

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

@@ -28,26 +28,27 @@
 @FunctionalInterface
 public interface Selector {
     /** A selector that always returns all provided values. */
-    Selector IDENTITY = (visitor, shapes) -> shapes;
+    Selector IDENTITY = (model, visitor, shapes) -> shapes;
 
     /**
      * Matches a selector to a set of shapes.
      *
+     * @param model Model used to resolve shapes with.
      * @param neighborProvider Provides neighbors for shapes.
      * @param shapes Matching context of shapes.
      * @return Returns the matching shapes.
      */
-    Set<Shape> select(NeighborProvider neighborProvider, Set<Shape> shapes);
+    Set<Shape> select(Model model, NeighborProvider neighborProvider, Set<Shape> shapes);
 
     /**
      * Matches a selector against a model using a custom neighbor visitor.
      *
-     * @param neighborProvider Provides neighbors for shapes
      * @param model Model to query.
+     * @param neighborProvider Provides neighbors for shapes
      * @return Returns the matching shapes.
      */
-    default Set<Shape> select(NeighborProvider neighborProvider, Model model) {
-        return select(neighborProvider, model.toSet());
+    default Set<Shape> select(Model model, NeighborProvider neighborProvider) {
+        return select(model, neighborProvider, model.toSet());
     }
 
     /**
@@ -57,7 +58,7 @@ default Set<Shape> select(NeighborProvider neighborProvider, Model model) {
      * @return Returns the matching shapes.
      */
     default Set<Shape> select(Model model) {
-        return select(NeighborProvider.of(model), model);
+        return select(model, NeighborProvider.of(model));
     }
 
     /**

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

@@ -17,6 +17,7 @@
 
 import java.util.Set;
 import java.util.stream.Collectors;
+import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.neighbor.NeighborProvider;
 import software.amazon.smithy.model.shapes.Shape;
 
@@ -28,7 +29,7 @@ final class ShapeTypeCategorySelector implements Selector {
     }
 
     @Override
-    public Set<Shape> select(NeighborProvider neighborProvider, Set<Shape> shapes) {
+    public Set<Shape> select(Model model, NeighborProvider neighborProvider, Set<Shape> shapes) {
         return shapes.stream().filter(shapeCategory::isInstance).collect(Collectors.toSet());
     }
 }