Fix dedicated io transform leaving unused shapes

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 {}