Fix dedicated io transform leaving unused shapes (#1419)

This commit changes the way `CreateDedicatedInputOutput` transform checks for
singular references to a shape being used as input or output in an operation.
It now checks to make sure the shape is only referenced by the specified
operation. This fixes an edge case where a shape was being used as both the
input and output for an operation, and would be left unused in the model
after the transformation was applied. Two new test cases were added to
verify the behavior when a shape is used as both input and output.

Fixes: #1373

smithy-model/src/main/java/software/amazon/smithy/model/transform/CreateDedicatedInputAndOutput.java

@@ -22,9 +22,7 @@
 import software.amazon.smithy.model.knowledge.NeighborProviderIndex;
 import software.amazon.smithy.model.knowledge.OperationIndex;
 import software.amazon.smithy.model.neighbor.NeighborProvider;
-import software.amazon.smithy.model.neighbor.Relationship;
 import software.amazon.smithy.model.neighbor.RelationshipDirection;
-import software.amazon.smithy.model.neighbor.RelationshipType;
 import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
@@ -60,41 +58,46 @@ Model transform(ModelTransformer transformer, Model model) {
             StructureShape output = operationIndex.expectOutputShape(operation);
             StructureShape updatedOutput = createdUpdatedOutput(model, operation, output, reverse);
 
-            if (!updatedInput.equals(input) || !updatedOutput.equals(output)) {
-                OperationShape.Builder builder = operation.toBuilder();
-                if (!updatedInput.equals(input)) {
-                    LOGGER.fine(() -> String.format("Updating operation input of %s from %s to %s",
-                                                    operation.getId(), input.getId(), updatedInput.getId()));
-                    updates.add(updatedInput);
-                    builder.input(updatedInput);
-                    // If the ID changed and the original is no longer referenced, then remove it.
-                    if (!input.getId().equals(updatedInput.getId())
-                            && isSingularReference(reverse, input, RelationshipType.INPUT)) {
-                        toRemove.add(input);
-                        LOGGER.info("Removing now unused input shape " + input.getId());
-                    }
+            boolean inputChanged = !input.equals(updatedInput);
+            boolean outputChanged = !output.equals(updatedOutput);
+
+            if (!inputChanged && !outputChanged) {
+                continue;
+            }
+
+            OperationShape.Builder builder = operation.toBuilder();
+            if (inputChanged) {
+                LOGGER.fine(() -> String.format("Updating operation input of %s from %s to %s",
+                        operation.getId(), input.getId(), updatedInput.getId()));
+                updates.add(updatedInput);
+                builder.input(updatedInput);
+                // If the ID changed and the original is no longer referenced, then remove it.
+                boolean idChanged = !input.getId().equals(updatedInput.getId());
+                if (idChanged && isSingularReference(reverse, input, operation)) {
+                    toRemove.add(input);
+                    LOGGER.info("Removing now unused input shape " + input.getId());
                 }
-                if (!updatedOutput.equals(output)) {
-                    LOGGER.fine(() -> String.format("Updating operation output of %s from %s to %s",
-                                                    operation.getId(), output.getId(), updatedOutput.getId()));
-                    updates.add(updatedOutput);
-                    builder.output(updatedOutput);
-                    // If the ID changed and the original is no longer referenced, then remove it.
-                    if (!output.getId().equals(updatedOutput.getId())
-                            && isSingularReference(reverse, output, RelationshipType.OUTPUT)) {
-                        toRemove.add(output);
-                        LOGGER.info("Removing now unused output shape " + output.getId());
-                    }
+            }
+            if (outputChanged) {
+                LOGGER.fine(() -> String.format("Updating operation output of %s from %s to %s",
+                        operation.getId(), output.getId(), updatedOutput.getId()));
+                updates.add(updatedOutput);
+                builder.output(updatedOutput);
+                // If the ID changed and the original is no longer referenced, then remove it.
+                boolean idChanged = !output.getId().equals(updatedOutput.getId());
+                if (idChanged && isSingularReference(reverse, output, operation)) {
+                    toRemove.add(output);
+                    LOGGER.info("Removing now unused output shape " + output.getId());
                 }
-                updates.add(builder.build());
             }
+            updates.add(builder.build());
         }
 
-        // Replace the operations and add new shapes.
-        Model result = transformer.replaceShapes(model, updates);
-
         // Remove no longer referenced shapes.
-        return transformer.removeShapes(result, toRemove);
+        Model result = transformer.removeShapes(model, toRemove);
+
+        // Replace the operations and add new shapes.
+        return transformer.replaceShapes(result, updates);
     }
 
     private StructureShape createdUpdatedInput(
@@ -105,7 +108,7 @@ private StructureShape createdUpdatedInput(
     ) {
         if (input.hasTrait(InputTrait.class)) {
             return renameShapeIfNeeded(model, input, operation, inputSuffix);
-        } else if (isDedicatedHeuristic(operation, input, reverse, RelationshipType.INPUT)) {
+        } else if (isDedicatedHeuristic(operation, input, reverse)) {
             LOGGER.fine(() -> "Attaching the @input trait to " + input.getId());
             InputTrait trait = new InputTrait(input.getSourceLocation());
             return renameShapeIfNeeded(model, input.toBuilder().addTrait(trait).build(), operation, inputSuffix);
@@ -122,7 +125,7 @@ private StructureShape createdUpdatedOutput(
     ) {
         if (output.hasTrait(OutputTrait.class)) {
             return renameShapeIfNeeded(model, output, operation, outputSuffix);
-        } else if (isDedicatedHeuristic(operation, output, reverse, RelationshipType.OUTPUT)) {
+        } else if (isDedicatedHeuristic(operation, output, reverse)) {
             LOGGER.fine(() -> "Attaching the @output trait to " + output.getId());
             OutputTrait trait = new OutputTrait(output.getSourceLocation());
             return renameShapeIfNeeded(model, output.toBuilder().addTrait(trait).build(), operation, outputSuffix);
@@ -153,36 +156,22 @@ private StructureShape renameShapeIfNeeded(
                 .build();
     }
 
-    private boolean isDedicatedHeuristic(
-            OperationShape operation,
-            StructureShape struct,
-            NeighborProvider reverse,
-            RelationshipType expected
-    ) {
+    private boolean isDedicatedHeuristic(OperationShape operation, StructureShape struct, NeighborProvider reverse) {
         // Only assume that a shape is dedicated to the operation its name starts with the operation name.
         if (!struct.getId().getName().startsWith(operation.getId().getName())) {
             return false;
         }
 
         // Check if the shape is only referenced as input or output.
-        return isSingularReference(reverse, struct, expected);
+        return isSingularReference(reverse, struct, operation);
     }
 
-    private boolean isSingularReference(NeighborProvider reverse, Shape shape, RelationshipType expected) {
-        int totalDirectedEdges = 0;
-
+    private boolean isSingularReference(NeighborProvider reverse, Shape shape, Shape expectedReferencingShape) {
         // We need to exclude inverted edges like MEMBER_CONTAINER, and only look for directed
-        // edges like INPUT and OUTPUT.
-        for (Relationship rel : reverse.getNeighbors(shape)) {
-            if (rel.getRelationshipType().getDirection() == RelationshipDirection.DIRECTED) {
-                totalDirectedEdges++;
-                if (rel.getRelationshipType() != expected) {
-                    return false;
-                }
-            }
-        }
-
-        return totalDirectedEdges == 1;
+        // edges to the expected shape.
+        return reverse.getNeighbors(shape).stream()
+                .filter(relationship -> relationship.getDirection() == RelationshipDirection.DIRECTED)
+                .allMatch(relationship -> relationship.getShape().equals(expectedReferencingShape));
     }
 
     private static StructureShape createSyntheticShape(

smithy-model/src/test/java/software/amazon/smithy/model/transform/CreateDedicatedInputAndOutputTest.java

@@ -169,6 +169,46 @@ public void handlesUnitTypes() {
                    Matchers.equalTo(ShapeId.from("smithy.api#Unit")));
     }
 
+    @Test
+    public void removesDisconnectedSharedShape() {
+        Model result = compareTransform("removes-disconnected-shared-shape", model ->
+                ModelTransformer.create().createDedicatedInputAndOutput(model, "Input", "Output")
+        );
+
+        assertThat(result.getShapeIds(), Matchers.not(Matchers.hasItem(ShapeId.from("smithy.example#MyGetFooOutput"))));
+
+        result.expectShape(ShapeId.from("smithy.example#GetFooInput")).expectTrait(InputTrait.class);
+        result.expectShape(ShapeId.from("smithy.example#GetFooOutput")).expectTrait(OutputTrait.class);
+
+        assertThat(result.expectShape(ShapeId.from("smithy.example#GetFooInput"))
+                            .expectTrait(OriginalShapeIdTrait.class)
+                            .getOriginalId(),
+                Matchers.equalTo(ShapeId.from("smithy.example#MyGetFooOutput")));
+
+        assertThat(result.expectShape(ShapeId.from("smithy.example#GetFooOutput"))
+                        .expectTrait(OriginalShapeIdTrait.class)
+                        .getOriginalId(),
+                Matchers.equalTo(ShapeId.from("smithy.example#MyGetFooOutput")));
+    }
+
+    @Test
+    public void createsDedicatedHeuristicForSharedShape() {
+        Model result = compareTransform("creates-dedicated-heuristic-for-shared", model ->
+                ModelTransformer.create().createDedicatedInputAndOutput(model, "Input", "Output")
+        );
+
+        result.expectShape(ShapeId.from("smithy.example#GetFooInput")).expectTrait(InputTrait.class);
+        result.expectShape(ShapeId.from("smithy.example#GetFooOutput")).expectTrait(OutputTrait.class);
+
+        assertThat(result.expectShape(ShapeId.from("smithy.example#GetFooInput"))
+                            .expectTrait(OriginalShapeIdTrait.class)
+                            .getOriginalId(),
+                Matchers.equalTo(ShapeId.from("smithy.example#GetFooOutput")));
+
+        assertThat(result.expectShape(ShapeId.from("smithy.example#GetFooOutput"))
+                        .getTrait(OriginalShapeIdTrait.class), Matchers.equalTo(Optional.empty()));
+    }
+
     @Test
     public void detectsConflictResolverLoop() {
         Assertions.assertThrows(ModelTransformException.class, () -> {

smithy-model/src/test/resources/software/amazon/smithy/model/transform/dedicated-input-output/creates-dedicated-heuristic-for-shared-after.smithy

@@ -0,0 +1,18 @@
+$version: "2.0"
+
+namespace smithy.example
+
+// `@output` is added to `GetFooOutput`
+// and a dedicated input shape is added
+operation GetFoo {
+    input: GetFooInput
+    output: GetFooOutput
+}
+
+@input
+structure GetFooInput {
+}
+
+@output
+structure GetFooOutput {
+}

smithy-model/src/test/resources/software/amazon/smithy/model/transform/dedicated-input-output/creates-dedicated-heuristic-for-shared-before.smithy

@@ -0,0 +1,12 @@
+$version: "2.0"
+
+namespace smithy.example
+
+// `GetFooOutput` is shared by input and output
+// and has the right naming for an output shape
+operation GetFoo {
+    input: GetFooOutput
+    output: GetFooOutput
+}
+
+structure GetFooOutput {}

smithy-model/src/test/resources/software/amazon/smithy/model/transform/dedicated-input-output/removes-disconnected-shared-shape-after.smithy

@@ -0,0 +1,16 @@
+$version: "2.0"
+
+namespace smithy.example
+
+// Dedicated input and output are created
+// and the old shared shape removed
+operation GetFoo {
+    input: GetFooInput
+    output: GetFooOutput
+}
+
+@input
+structure GetFooInput {}
+
+@output
+structure GetFooOutput {}

smithy-model/src/test/resources/software/amazon/smithy/model/transform/dedicated-input-output/removes-disconnected-shared-shape-before.smithy

@@ -0,0 +1,12 @@
+$version: "2.0"
+
+namespace smithy.example
+
+// A shape is shared by both input and output
+// but is not named properly for an output
+operation GetFoo {
+    input: MyGetFooOutput
+    output: MyGetFooOutput
+}
+
+structure MyGetFooOutput {}