smithy-lang · JordonPhillips · Dec 11, 2023 · Dec 8, 2023
@@ -18,6 +18,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -37,6 +38,7 @@
 import software.amazon.smithy.model.traits.OutputTrait;
 import software.amazon.smithy.model.traits.UnitTypeTrait;
 import software.amazon.smithy.utils.ListUtils;
+import software.amazon.smithy.utils.SetUtils;
 
 /**
  * Index of operation IDs to their resolved input, output, and error
@@ -50,11 +52,20 @@ public final class OperationIndex implements KnowledgeIndex {
     private final Map<ShapeId, StructureShape> inputs = new HashMap<>();
     private final Map<ShapeId, StructureShape> outputs = new HashMap<>();
     private final Map<ShapeId, List<StructureShape>> errors = new HashMap<>();
+    private final Map<ShapeId, Set<OperationShape>> boundInputOperations = new HashMap<>();
+    private final Map<ShapeId, Set<OperationShape>> boundOutputOperations = new HashMap<>();
+    private final Map<ShapeId, Set<Shape>> boundErrorShapes = new HashMap<>();
 
     public OperationIndex(Model model) {
         for (OperationShape operation : model.getOperationShapes()) {
-            getStructure(model, operation.getInputShape()).ifPresent(shape -> inputs.put(operation.getId(), shape));
-            getStructure(model, operation.getOutputShape()).ifPresent(shape -> outputs.put(operation.getId(), shape));
+            getStructure(model, operation.getInputShape()).ifPresent(shape -> {
+                inputs.put(operation.getId(), shape);
+                boundInputOperations.computeIfAbsent(shape.getId(), id -> new HashSet<>()).add(operation);
+            });
+            getStructure(model, operation.getOutputShape()).ifPresent(shape -> {
+                outputs.put(operation.getId(), shape);
+                boundOutputOperations.computeIfAbsent(shape.getId(), id -> new HashSet<>()).add(operation);
+            });
             addErrorsFromShape(model, operation.getId(), operation.getErrors());
         }
 
@@ -65,8 +76,10 @@ public OperationIndex(Model model) {
 
     private void addErrorsFromShape(Model model, ShapeId source, List<ShapeId> errorShapeIds) {
         List<StructureShape> errorShapes = new ArrayList<>(errorShapeIds.size());
+        Shape sourceShape = model.expectShape(source);
         for (ShapeId target : errorShapeIds) {
             model.getShape(target).flatMap(Shape::asStructureShape).ifPresent(errorShapes::add);
+            boundErrorShapes.computeIfAbsent(target, id -> new HashSet<>()).add(sourceShape);
         }
         errors.put(source, errorShapes);
     }
@@ -158,6 +171,16 @@ public boolean isInputStructure(ToShapeId structureId) {
         return false;
     }
 
+    /**
+     * Gets all the operations that bind the given shape as input.
+     *
+     * @param input The structure that may be used as input.
+     * @return Returns a set of operations that bind the given input shape.
+     */
+    public Set<OperationShape> getInputBindings(ToShapeId input) {
+        return SetUtils.copyOf(boundInputOperations.getOrDefault(input.toShapeId(), Collections.emptySet()));
+    }
+
     /**
      * Gets the optional output structure of an operation, and returns an
      * empty optional if the output targets {@code smithy.api#Unit}.
@@ -241,6 +264,16 @@ public boolean isOutputStructure(ToShapeId structureId) {
         return false;
     }
 
+    /**
+     * Gets all the operations that bind the given shape as output.
+     *
+     * @param output The structure that may be used as output.
+     * @return Returns a set of operations that bind the given output shape.
+     */
+    public Set<OperationShape> getOutputBindings(ToShapeId output) {
+        return SetUtils.copyOf(boundOutputOperations.getOrDefault(output.toShapeId(), Collections.emptySet()));
+    }
+
     /**
      * Gets the list of error structures defined on an operation.
      *
@@ -275,4 +308,14 @@ public List<StructureShape> getErrors(ToShapeId service, ToShapeId operation) {
     private Optional<StructureShape> getStructure(Model model, ToShapeId id) {
         return model.getShape(id.toShapeId()).flatMap(Shape::asStructureShape);
     }
+
+    /**
+     * Gets all the operations and services that bind the given shape as an error.
+     *
+     * @param error The structure that may be used as an error.
+     * @return Returns a set of operations and services that bind the given error shape.
+     */
+    public Set<Shape> getErrorBindings(ToShapeId error) {
+        return SetUtils.copyOf(boundErrorShapes.getOrDefault(error.toShapeId(), Collections.emptySet()));
+    }
 }
@@ -23,15 +23,18 @@
 
 import java.util.Collections;
 import java.util.Optional;
+import java.util.Set;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.StructureShape;
 import software.amazon.smithy.model.traits.UnitTypeTrait;
+import software.amazon.smithy.utils.SetUtils;
 
 public class OperationIndexTest {
 
@@ -40,7 +43,7 @@ public class OperationIndexTest {
     @BeforeAll
     public static void before() {
         model = Model.assembler()
-                .addImport(OperationIndexTest.class.getResource("operation-index-test.json"))
+                .addImport(OperationIndexTest.class.getResource("operation-index-test.smithy"))
                 .assemble()
                 .unwrap();
     }
@@ -103,6 +106,18 @@ public void determinesIfShapeIsUsedAsInput() {
         assertThat(opIndex.isInputStructure(output), is(false));
     }
 
+    @Test
+    public void getsInputBindings() {
+        OperationIndex index = OperationIndex.of(model);
+        Set<OperationShape> actual = index.getInputBindings(ShapeId.from("ns.foo#Input"));
+        Set<OperationShape> expected = SetUtils.of(
+                model.expectShape(ShapeId.from("ns.foo#B"), OperationShape.class),
+                model.expectShape(ShapeId.from("ns.foo#C"), OperationShape.class)
+        );
+        assertThat(actual, equalTo(expected));
+        assertThat(index.getInputBindings(ShapeId.from("ns.foo#Output")), empty());
+    }
+
     @Test
     public void determinesIfShapeIsUsedAsOutput() {
         OperationIndex opIndex = OperationIndex.of(model);
@@ -113,6 +128,18 @@ public void determinesIfShapeIsUsedAsOutput() {
         assertThat(opIndex.isOutputStructure(input), is(false));
     }
 
+    @Test
+    public void getsOutputBindings() {
+        OperationIndex index = OperationIndex.of(model);
+        Set<OperationShape> actual = index.getOutputBindings(ShapeId.from("ns.foo#Output"));
+        Set<OperationShape> expected = SetUtils.of(
+                model.expectShape(ShapeId.from("ns.foo#B"), OperationShape.class),
+                model.expectShape(ShapeId.from("ns.foo#C"), OperationShape.class)
+        );
+        assertThat(actual, equalTo(expected));
+        assertThat(index.getOutputBindings(ShapeId.from("ns.foo#Input")), empty());
+    }
+
     @Test
     public void getsOperationErrorsAndInheritedErrors() {
         OperationIndex opIndex = OperationIndex.of(model);
@@ -130,4 +157,21 @@ public void getsOperationErrorsAndInheritedErrors() {
         assertThat(opIndex.getErrors(b), containsInAnyOrder(error1, error2));
         assertThat(opIndex.getErrors(service, b), containsInAnyOrder(error1, error2, common1, common2));
     }
+
+    @Test
+    public void getsErrorBindings() {
+        OperationIndex index = OperationIndex.of(model);
+        Set<Shape> actual = index.getErrorBindings(ShapeId.from("ns.foo#CommonError1"));
+        Set<Shape> expected = SetUtils.of(
+                model.expectShape(ShapeId.from("ns.foo#MyService"), ServiceShape.class),
+                model.expectShape(ShapeId.from("ns.foo#C"), OperationShape.class)
+        );
+        assertThat(actual, equalTo(expected));
+
+        actual = index.getErrorBindings(ShapeId.from("ns.foo#Error1"));
+        expected = SetUtils.of(model.expectShape(ShapeId.from("ns.foo#B"), OperationShape.class));
+        assertThat(actual, equalTo(expected));
+
+        assertThat(index.getOutputBindings(ShapeId.from("ns.foo#UnusedError")), empty());
+    }
 }
@@ -0,0 +1,58 @@
+$version: "2.0"
+
+namespace ns.foo
+
+service MyService {
+    version: "2017-01-17"
+    operations: [
+        A
+        B
+    ]
+    errors: [
+        CommonError1
+        CommonError2
+    ]
+}
+
+@readonly
+operation A {
+    input: Unit
+    output: Unit
+}
+
+@readonly
+operation B {
+    input: Input
+    output: Output
+    errors: [
+        Error1
+        Error2
+    ]
+}
+
+operation C {
+    input: Input
+    output: Output,
+    errors: [
+        CommonError1
+    ]
+}
+
+@error("server")
+structure CommonError1 {}
+
+@error("server")
+structure CommonError2 {}
+
+@error("client")
+structure Error1 {}
+
+@error("server")
+structure Error2 {}
+
+@error("client")
+structure UnusedError {}
+
+structure Input {}
+
+structure Output {}