Refactor to create an ID_REF relationship type

Previous impl wouldn't consider shapes referenced through multiple
levels of idRef as connected, since idRef'd shapes could themselves
be connected to other shapes that have idRef references. Refactoring
to a custom NeighborProvider means the shape walker just traverses
those relationships in a single pass.

smithy-model/src/main/java/software/amazon/smithy/model/knowledge/NeighborProviderIndex.java

@@ -32,6 +32,7 @@ public final class NeighborProviderIndex implements KnowledgeIndex {
     private volatile NeighborProvider reversed;
     private volatile NeighborProvider providerWithTraits;
     private volatile NeighborProvider reversedWithTraits;
+    private volatile NeighborProvider providerWithIdRefs;
 
     public NeighborProviderIndex(Model model) {
         provider = NeighborProvider.precomputed(model);
@@ -76,6 +77,28 @@ public NeighborProvider getProviderWithTraitRelationships() {
         return result;
     }
 
+    /**
+     * Gets the neighbor provider that includes idRef relationships.
+     *
+     * @return Returns the provider.
+     */
+    public NeighborProvider getProviderWithIdRefRelationships() {
+        NeighborProvider result = providerWithIdRefs;
+
+        if (result == null) {
+            Model model = getOrThrowModel();
+            synchronized (this) {
+                result = providerWithIdRefs;
+                if (result == null) {
+                    providerWithIdRefs = result = NeighborProvider.cached(
+                            NeighborProvider.withIdRefRelationships(model, provider));
+                }
+            }
+        }
+
+        return result;
+    }
+
     /**
      * Gets a reversed, bottom up neighbor provider.
      *

smithy-model/src/main/java/software/amazon/smithy/model/neighbor/IdRefShapeReferences.java → smithy-model/src/main/java/software/amazon/smithy/model/neighbor/IdRefShapeRelationships.java

@@ -5,8 +5,10 @@
 
 package software.amazon.smithy.model.neighbor;
 
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.node.Node;
@@ -21,55 +23,55 @@
 import software.amazon.smithy.model.traits.TraitDefinition;
 
 /**
- * Finds all shapes referenced by {@link IdRefTrait} from within trait
- * values.
+ * Finds all {@link RelationshipType#ID_REF} relationships in a model.
  *
  * <p>This works by finding all paths from {@link TraitDefinition} shapes
  * to shapes with {@link IdRefTrait}, then using those paths to search
  * the node values of each application of the trait to extract the {@link ShapeId}
  * value. Because we don't have a fixed set of traits known to potentially have
  * idRef members, this has to be done dynamically.
  */
-final class IdRefShapeReferences {
+final class IdRefShapeRelationships {
     private static final Selector WITH_ID_REF = Selector.parse("[trait|idRef]");
     private final Model model;
-    private final Set<ShapeId> found = new HashSet<>();
+    private final Map<ShapeId, Set<Relationship>> relationships = new HashMap<>();
 
-    IdRefShapeReferences(Model model) {
+    IdRefShapeRelationships(Model model) {
         this.model = model;
     }
 
-    Set<ShapeId> compute(Set<Shape> connected) {
+    Map<ShapeId, Set<Relationship>> getRelationships() {
         PathFinder finder = PathFinder.create(model);
         for (Shape traitDef : model.getShapesWithTrait(TraitDefinition.class)) {
             if (traitDef.hasTrait(IdRefTrait.class)) {
                 // PathFinder doesn't handle the case where the trait def has the idRef
                 NodeQuery query = new NodeQuery().self();
-                addReferences(traitDef, query, connected);
+                addRelationships(traitDef, query);
                 continue;
             }
             List<PathFinder.Path> paths = finder.search(traitDef, WITH_ID_REF);
             if (!paths.isEmpty()) {
                 for (PathFinder.Path path : paths) {
                     NodeQuery query = buildNodeQuery(path);
-                    addReferences(traitDef, query, connected);
+                    addRelationships(traitDef, query);
                 }
             }
         }
-        return found;
+        return relationships;
     }
 
-    private void addReferences(Shape traitDef, NodeQuery query, Set<Shape> connected) {
-        model.getShapesWithTrait(traitDef.getId()).stream()
-                .filter(connected::contains)
-                .forEach(shape -> {
-                    Trait trait = shape.findTrait(traitDef.getId()).get(); // We already know the shape has the trait.
-                    Node node = trait.toNode();
-                    query.execute(node).forEach(n ->
-                            // Invalid shape ids are handled by the idRef trait validator, so ignore them here.
-                            n.asStringNode().flatMap(StringNode::asShapeId).ifPresent(found::add)
-                    );
-                });
+    private void addRelationships(Shape traitDef, NodeQuery query) {
+        model.getShapesWithTrait(traitDef.getId()).forEach(shape -> {
+            Trait trait = shape.findTrait(traitDef.getId()).get(); // We already know the shape has the trait.
+            Node node = trait.toNode();
+            // Invalid shape ids are handled by the idRef trait validator, so ignore them here.
+            query.execute(node).forEach(n -> n.asStringNode()
+                    .flatMap(StringNode::asShapeId)
+                    .flatMap(model::getShape)
+                    .map(referenced -> Relationship.create(shape, RelationshipType.ID_REF, referenced))
+                    .ifPresent(rel -> relationships
+                            .computeIfAbsent(rel.getShape().getId(), id -> new HashSet<>()).add(rel)));
+        });
     }
 
     private static NodeQuery buildNodeQuery(PathFinder.Path path) {

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

@@ -26,6 +26,7 @@
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.traits.IdRefTrait;
 import software.amazon.smithy.utils.ListUtils;
 
 /**
@@ -91,6 +92,29 @@ static NeighborProvider withTraitRelationships(Model model, NeighborProvider nei
         };
     }
 
+    /**
+     * Creates a NeighborProvider that includes {@link RelationshipType#ID_REF}
+     * relationships.
+     *
+     * @param model Model to use to look up shapes referenced by {@link IdRefTrait}.
+     * @param neighborProvider Provider to wrap.
+     * @return Returns the created neighbor provider.
+     */
+    static NeighborProvider withIdRefRelationships(Model model, NeighborProvider neighborProvider) {
+        Map<ShapeId, Set<Relationship>> idRefRelationships = new IdRefShapeRelationships(model).getRelationships();
+        return shape -> {
+            List<Relationship> relationships = neighborProvider.getNeighbors(shape);
+
+            if (!idRefRelationships.containsKey(shape.getId())) {
+                return relationships;
+            }
+
+            relationships = new ArrayList<>(relationships);
+            relationships.addAll(idRefRelationships.get(shape.getId()));
+            return relationships;
+        };
+    }
+
     /**
      * Creates a NeighborProvider that precomputes the neighbors of a model.
      *

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

@@ -26,8 +26,10 @@
 import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.ResourceShape;
 import software.amazon.smithy.model.shapes.SetShape;
+import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.StructureShape;
 import software.amazon.smithy.model.shapes.UnionShape;
+import software.amazon.smithy.model.traits.IdRefTrait;
 import software.amazon.smithy.model.traits.TraitDefinition;
 import software.amazon.smithy.utils.SmithyInternalApi;
 
@@ -214,7 +216,37 @@ public enum RelationshipType {
      * Relationship that exists between a structure or union and a mixin applied
      * to the shape.
      */
-    MIXIN("mixin", RelationshipDirection.DIRECTED);
+    MIXIN("mixin", RelationshipDirection.DIRECTED),
+
+    /**
+     * Relationships that exist between a shape and another shape referenced by an
+     * {@link IdRefTrait}.
+     *
+     * <p>This relationship is formed by applying a trait with a value containing a
+     * reference to another {@link ShapeId}. For
+     * example:
+     * <pre>
+     * {@code
+     * @trait
+     * structure myRef {
+     *     @idRef
+     *     shape: String
+     * }
+     *
+     * // @myRef trait applied, and the value references the shape `Referenced`
+     * @myRef(shape: Referenced)
+     * structure WithMyRef {}
+     *
+     * string Referenced
+     * }
+     * </pre>
+     *
+     * <p>This kind of relationship is not returned by default from a
+     * {@link NeighborProvider}. You must explicitly wrap a {@link NeighborProvider}
+     * with {@link NeighborProvider#withIdRefRelationships(Model, NeighborProvider)}
+     * in order to yield idRef relationships.
+     */
+    ID_REF(null, RelationshipDirection.DIRECTED);
 
     private String selectorLabel;
     private RelationshipDirection direction;

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

@@ -22,13 +22,14 @@
 import software.amazon.smithy.model.knowledge.NeighborProviderIndex;
 import software.amazon.smithy.model.loader.Prelude;
 import software.amazon.smithy.model.shapes.Shape;
-import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.TraitDefinition;
 import software.amazon.smithy.utils.FunctionalUtils;
 
 /**
  * Finds shapes that are not connected to a service shape, are not trait
- * definitions, and are not referenced by trait definitions.
+ * definitions, are not referenced by trait definitions, and are not
+ * referenced in trait values through
+ * {@link software.amazon.smithy.model.traits.IdRefTrait}.
  *
  * <p>Prelude shapes are never considered unreferenced.
  */
@@ -54,7 +55,7 @@ public UnreferencedShapes(Predicate<Shape> keepFilter) {
      * @return Returns the unreferenced shapes.
      */
     public Set<Shape> compute(Model model) {
-        Walker shapeWalker = new Walker(NeighborProviderIndex.of(model).getProvider());
+        Walker shapeWalker = new Walker(NeighborProviderIndex.of(model).getProviderWithIdRefRelationships());
 
         // Find all shapes connected to any service shape.
         Set<Shape> connected = new HashSet<>();
@@ -67,12 +68,6 @@ public Set<Shape> compute(Model model) {
             connected.addAll(shapeWalker.walkShapes(trait));
         }
 
-        for (ShapeId referencedThroughIdRef : new IdRefShapeReferences(model).compute(connected)) {
-            // Referenced shapes may not exist in the model, but we don't want to throw.
-            model.getShape(referencedThroughIdRef)
-                    .ifPresent(shape -> connected.addAll(shapeWalker.walkShapes(shape)));
-        }
-
         // Any shape that wasn't identified as connected to a service is considered unreferenced.
         Set<Shape> result = new HashSet<>();
         for (Shape shape : model.toSet()) {

smithy-model/src/test/java/software/amazon/smithy/model/neighbor/NeighborProviderTest.java

@@ -1,13 +1,18 @@
 package software.amazon.smithy.model.neighbor;
 
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
 
 import java.util.List;
+import java.util.stream.Collectors;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.StringShape;
 import software.amazon.smithy.model.traits.SensitiveTrait;
@@ -42,4 +47,95 @@ public void canGetTraitRelationshipsFromShapeWithNoTraits() {
 
         assertThat(relationships, empty());
     }
+
+    @ParameterizedTest
+    @CsvSource({
+            "One,Ref1",
+            "Two,Ref2",
+            "Three,Ref3",
+            "Four,Ref4",
+            "Five,Ref5",
+            "Six,Ref6",
+            "Seven,Ref7",
+            "Eight,Ref8",
+            "Nine,Ref9",
+            "Ten,Ref10",
+            "Eleven,Ref11",
+            "Twelve,Ref12",
+            "Thirteen,Ref13",
+            "Fourteen,Ref14"
+    })
+    public void canGetIdRefRelationships(String shapeName, String referencedShapeName) {
+        Model model = Model.assembler()
+                .addImport(getClass().getResource("idref-neighbors.smithy"))
+                .assemble()
+                .unwrap();
+        NeighborProvider provider = NeighborProvider.of(model);
+        provider = NeighborProvider.withIdRefRelationships(model, provider);
+
+        Shape shape = model.expectShape(ShapeId.fromParts("com.foo", shapeName));
+        Shape ref = model.expectShape(ShapeId.fromParts("com.foo", referencedShapeName));
+        List<Relationship> relationships = provider.getNeighbors(shape).stream()
+                .filter(relationship -> relationship.getRelationshipType().equals(RelationshipType.ID_REF))
+                .collect(Collectors.toList());
+
+        assertThat(relationships, containsInAnyOrder(
+                equalTo(Relationship.create(shape, RelationshipType.ID_REF, ref))));
+    }
+
+    @Test
+    public void canGetIdRefRelationshipsThroughTraitDefs() {
+        Model model = Model.assembler()
+                .addImport(getClass().getResource("idref-neighbors-in-trait-def.smithy"))
+                .assemble()
+                .unwrap();
+        NeighborProvider provider = NeighborProvider.of(model);
+        provider = NeighborProvider.withIdRefRelationships(model, provider);
+
+        Shape shape = model.expectShape(ShapeId.from("com.foo#WithRefStructTrait"));
+        Shape ref = model.expectShape(ShapeId.from("com.foo#OtherReferenced"));
+        List<Relationship> relationships = provider.getNeighbors(shape).stream()
+                .filter(relationship -> relationship.getRelationshipType().equals(RelationshipType.ID_REF))
+                .collect(Collectors.toList());
+
+        assertThat(relationships, containsInAnyOrder(
+                equalTo(Relationship.create(shape, RelationshipType.ID_REF, ref))));
+
+        Shape shape1 = model.expectShape(ShapeId.from("com.foo#refStruct$other"));
+        Shape ref1 = model.expectShape(ShapeId.from("com.foo#ReferencedInTraitDef"));
+        List<Relationship> relationships1 = provider.getNeighbors(shape1).stream()
+                .filter(relationship -> relationship.getRelationshipType().equals(RelationshipType.ID_REF))
+                .collect(Collectors.toList());
+
+        assertThat(relationships1, containsInAnyOrder(
+                equalTo(Relationship.create(shape1, RelationshipType.ID_REF, ref1))));
+    }
+
+    @Test
+    public void canGetIdRefRelationshipsThroughMultipleLevelsOfIdRef() {
+        Model model = Model.assembler()
+                .addImport(getClass().getResource("idref-neighbors-multiple-levels.smithy"))
+                .assemble()
+                .unwrap();
+        NeighborProvider provider = NeighborProvider.of(model);
+        provider = NeighborProvider.withIdRefRelationships(model, provider);
+
+        Shape shape = model.expectShape(ShapeId.from("com.foo#WithIdRef"));
+        Shape ref = model.expectShape(ShapeId.from("com.foo#Referenced"));
+        List<Relationship> relationships = provider.getNeighbors(shape).stream()
+                .filter(relationship -> relationship.getRelationshipType().equals(RelationshipType.ID_REF))
+                .collect(Collectors.toList());
+
+        assertThat(relationships, containsInAnyOrder(
+                Relationship.create(shape, RelationshipType.ID_REF, ref)));
+
+        Shape shape1 = model.expectShape(ShapeId.from("com.foo#ConnectedThroughReference"));
+        Shape ref1 = model.expectShape(ShapeId.from("com.foo#AnotherReferenced"));
+        List<Relationship> relationships1 = provider.getNeighbors(shape1).stream()
+                .filter(relationship -> relationship.getRelationshipType().equals(RelationshipType.ID_REF))
+                .collect(Collectors.toList());
+
+        assertThat(relationships1, containsInAnyOrder(
+                Relationship.create(shape1, RelationshipType.ID_REF, ref1)));
+    }
 }