Add back vendorParamsShape if removed by transformers

Adds a ModelTransformerPlugin to smithy-smoke-test-traits that will
add back shapes referenced by a smoke test case vendorParamsShape that
were removed by model transforms. Transforms like removeUnusedShapes
will remove these shapes since they aren't connected to the model graph
as @idref's don't create edges.

To accomplish this, ModelTransformerPlugin was given an `order` used to
sort the plugins. This was required to make the new plugin run after
all other plugins, which may remove more shapes that need to be added
back.

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

@@ -27,6 +27,7 @@
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.loader.ModelAssembler;
 import software.amazon.smithy.model.neighbor.UnreferencedShapes;
@@ -51,7 +52,10 @@ public final class ModelTransformer {
     private final List<ModelTransformerPlugin> plugins;
 
     private ModelTransformer(List<ModelTransformerPlugin> plugins) {
-        this.plugins = ListUtils.copyOf(plugins);
+        List<ModelTransformerPlugin> copy = ListUtils.copyOf(plugins).stream()
+                .sorted(Comparator.comparingInt(ModelTransformerPlugin::order))
+                .collect(Collectors.toList());
+        this.plugins = ListUtils.copyOf(copy);
     }
 
     /**

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

@@ -38,4 +38,17 @@ public interface ModelTransformerPlugin {
     default Model onRemove(ModelTransformer transformer, Collection<Shape> removed, Model model) {
         return model;
     }
+
+    /**
+     * Defines the sort order of the plugin, a value from -128 to 127.
+     *
+     * <p>Plugins are applied according to this sort order. Lower values
+     * are executed before higher values (for example, -128 comes before 0,
+     * 0 comes before 127). Plugins default to 0.
+     *
+     * @return Returns the sort order, defaulting to 0.
+     */
+    default byte order() {
+        return 0;
+    }
 }

smithy-smoke-test-traits/src/main/java/software/amazon/smithy/smoketests/traits/transform/KeepVendorParamsShapes.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.smoketests.traits.transform;
+
+import java.util.Collection;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.neighbor.Walker;
+import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.transform.ModelTransformer;
+import software.amazon.smithy.model.transform.ModelTransformerPlugin;
+import software.amazon.smithy.smoketests.traits.SmokeTestCase;
+import software.amazon.smithy.smoketests.traits.SmokeTestsTrait;
+
+/**
+ * Runs after all other {@link ModelTransformerPlugin}s, adding back any shapes referenced by a
+ * {@code SmokeTestCase.vendorParamsShape} and all connected shapes, that were removed by previous transforms.
+ *
+ * <p>Since these shapes are referenced from within trait
+ * values, they don't create an edge in the model graph. This means transforms like
+ * <a href="https://smithy.io/2.0/guides/building-models/build-config.html#removeunusedshapes">removeUnusedShapes</a>
+ * will remove vendor params shapes, causing {@link software.amazon.smithy.smoketests.traits.SmokeTestCaseValidator}
+ * to fail.
+ */
+public class KeepVendorParamsShapes implements ModelTransformerPlugin {
+    @Override
+    public byte order() {
+        // This plugin has to run last, in case previous plugins removed any of the vendor params shapes.
+        return Byte.MAX_VALUE;
+    }
+
+    @Override
+    public Model onRemove(ModelTransformer transformer, Collection<Shape> removed, Model model) {
+        Set<ShapeId> vendorParamsShapeIds = model.getShapesWithTrait(SmokeTestsTrait.class).stream()
+                .flatMap(shape -> shape.expectTrait(SmokeTestsTrait.class).getTestCases().stream())
+                .map(SmokeTestCase::getVendorParamsShape)
+                .filter(Optional::isPresent)
+                .map(Optional::get)
+                .collect(Collectors.toSet());
+
+        // Only consider vendor params shapes that were removed.
+        vendorParamsShapeIds.removeAll(model.getShapeIds());
+        if (vendorParamsShapeIds.isEmpty()) {
+            return model;
+        }
+
+        Model.Builder builder = model.toBuilder();
+
+        // Need to add back all the shapes connected to the vendor params shape as well.
+        Model removedShapesModel = Model.builder().addShapes(removed).build();
+        Walker removedShapesWalker = new Walker(removedShapesModel);
+        for (ShapeId removedVendorParamsShapeId : vendorParamsShapeIds) {
+            Shape removedShape = removedShapesModel.expectShape(removedVendorParamsShapeId);
+            Set<Shape> connected = removedShapesWalker.walkShapes(removedShape);
+            builder.addShapes(connected);
+        }
+
+        return builder.build();
+    }
+}

smithy-smoke-test-traits/src/main/resources/META-INF/services/software.amazon.smithy.model.transform.ModelTransformerPlugin

@@ -0,0 +1 @@
+software.amazon.smithy.smoketests.traits.transform.KeepVendorParamsShapes

smithy-smoke-test-traits/src/test/java/software/amazon/smithy/smoketests/traits/transform/KeepVendorParamsShapesTest.java

@@ -0,0 +1,106 @@
+package software.amazon.smithy.smoketests.traits.transform;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
+import static org.hamcrest.Matchers.equalTo;
+
+import java.util.Optional;
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.transform.ModelTransformer;
+import software.amazon.smithy.smoketests.traits.SmokeTestsTrait;
+import software.amazon.smithy.utils.ListUtils;
+
+public class KeepVendorParamsShapesTest {
+    @Test
+    public void keepsOnlyVendorParams() {
+        Model model = Model.assembler()
+                .discoverModels()
+                .addImport(getClass().getResource("vendor-params-and-other-unused-shapes.smithy"))
+                .assemble()
+                .unwrap();
+
+        ModelTransformer transformer = ModelTransformer.create();
+        Model transformed = transformer.removeUnreferencedShapes(model);
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#VendorParams")), not(equalTo(Optional.empty())));
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#VendorParams$foo")), not(equalTo(Optional.empty())));
+
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#Unused")), equalTo(Optional.empty()));
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#Unused$unusedMember")), equalTo(Optional.empty()));
+    }
+
+    @Test
+    public void doesntKeepVendorParamsOnUnconnectedOperations() {
+        Model model = Model.assembler()
+                .discoverModels()
+                .addImport(getClass().getResource("vendor-params-with-unconnected-operation.smithy"))
+                .assemble()
+                .unwrap();
+
+        Model transformed = ModelTransformer.create().removeUnreferencedShapes(model);
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#GetFoo")), equalTo(Optional.empty()));
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#VendorParams")), equalTo(Optional.empty()));
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#VendorParams$foo")), equalTo(Optional.empty()));
+    }
+
+    @Test
+    public void doesntKeepIfSmokeTestsAreRemoved() {
+        Model model = Model.assembler()
+                .discoverModels()
+                .addImport(getClass().getResource("vendor-params-and-other-unused-shapes.smithy"))
+                .assemble()
+                .unwrap();
+
+        ModelTransformer transformer = ModelTransformer.create();
+        Model transformed = transformer.removeUnreferencedShapes(transformer
+                .removeTraitsIf(model, (shape, trait) -> trait.toShapeId().equals(SmokeTestsTrait.ID)));
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#VendorParams")), equalTo(Optional.empty()));
+    }
+
+    @Test
+    public void keepsShapesReferencedByVendorParamsShape() {
+        Model model = Model.assembler()
+                .discoverModels()
+                .addImport(getClass().getResource("vendor-params-with-nested-shapes.smithy"))
+                .assemble()
+                .unwrap();
+
+        Model transformed = ModelTransformer.create().removeUnreferencedShapes(model);
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#VendorParams")), not(equalTo(Optional.empty())));
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#VendorParams$nestedStruct")), not(equalTo(Optional.empty())));
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#NestedStruct")), not(equalTo(Optional.empty())));
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#NestedStruct$nestedString")), not(equalTo(Optional.empty())));
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#NestedString")), not(equalTo(Optional.empty())));
+    }
+
+    @Test
+    public void doesntKeepShapesThatTargetVendorParams() {
+        Model model = Model.assembler()
+                .discoverModels()
+                .addImport(getClass().getResource("vendor-params-referenced-by-unconnected-shape.smithy"))
+                .assemble()
+                .unwrap();
+
+        Model transformed = ModelTransformer.create().removeUnreferencedShapes(model);
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#VendorParams")), not(equalTo(Optional.empty())));
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#Unconnected")), equalTo(Optional.empty()));
+        assertThat(transformed.getShape(ShapeId.from("smithy.example#Unconnected$vendorParams")), equalTo(Optional.empty()));
+    }
+
+    @Test
+    public void shapesConnectedToVendorParamsCanStillBeRemoved() {
+        // NOTE: Removing `NestedStruct` also removes members that target it, mutating `VendorParams`.
+        Model model = Model.assembler()
+                .discoverModels()
+                .addImport(getClass().getResource("vendor-params-with-nested-shapes.smithy"))
+                .assemble()
+                .unwrap();
+
+        ShapeId connected = ShapeId.from("smithy.example#NestedStruct");
+        Model removeConnected = ModelTransformer.create()
+                .removeShapes(model, ListUtils.of(model.expectShape(connected)));
+        assertThat(removeConnected.getShape(connected).isPresent(), is(false));
+    }
+}

smithy-smoke-test-traits/src/test/resources/software/amazon/smithy/smoketests/traits/transform/vendor-params-and-other-unused-shapes.smithy

@@ -0,0 +1,32 @@
+$version: "2.0"
+
+namespace smithy.example
+
+use smithy.test#smokeTests
+
+service HelloService {
+    version: "2024-01-17"
+    operations: [SayHello]
+}
+
+@smokeTests([
+    {
+        id: "with_vendor_params_shape",
+        expect: {
+            success: {}
+        },
+        vendorParams: {
+            foo: "Bar"
+        },
+        vendorParamsShape: VendorParams
+    }
+])
+operation SayHello {}
+
+structure VendorParams {
+    foo: String
+}
+
+structure Unused {
+    unusedMember: String
+}

smithy-smoke-test-traits/src/test/resources/software/amazon/smithy/smoketests/traits/transform/vendor-params-referenced-by-unconnected-shape.smithy

@@ -0,0 +1,32 @@
+$version: "2.0"
+
+namespace smithy.example
+
+use smithy.test#smokeTests
+
+service HelloService {
+    version: "2024-01-17"
+    operations: [SayHello]
+}
+
+@smokeTests([
+    {
+        id: "with_vendor_params_shape",
+        expect: {
+            success: {}
+        },
+        vendorParams: {
+            foo: "Bar"
+        },
+        vendorParamsShape: VendorParams
+    }
+])
+operation SayHello {}
+
+structure VendorParams {
+    foo: String
+}
+
+structure Unconnected {
+    vendorParams: VendorParams
+}

smithy-smoke-test-traits/src/test/resources/software/amazon/smithy/smoketests/traits/transform/vendor-params-with-nested-shapes.smithy

@@ -0,0 +1,36 @@
+$version: "2.0"
+
+namespace smithy.example
+
+use smithy.test#smokeTests
+
+service HelloService {
+    version: "2024-01-17"
+    operations: [SayHello]
+}
+
+@smokeTests([
+    {
+        id: "with_vendor_params_shape",
+        expect: {
+            success: {}
+        },
+        vendorParams: {
+            nestedStruct: {
+                nestedString: "foo"
+            }
+        },
+        vendorParamsShape: VendorParams
+    }
+])
+operation SayHello {}
+
+structure VendorParams {
+    nestedStruct: NestedStruct
+}
+
+structure NestedStruct {
+    nestedString: NestedString
+}
+
+string NestedString

smithy-smoke-test-traits/src/test/resources/software/amazon/smithy/smoketests/traits/transform/vendor-params-with-unconnected-operation.smithy

@@ -0,0 +1,27 @@
+$version: "2.0"
+
+namespace smithy.example
+
+use smithy.test#smokeTests
+
+service HelloService {
+    version: "2024-01-17"
+}
+
+@smokeTests([
+    {
+        id: "with_vendor_params_shape",
+        expect: {
+            success: {}
+        },
+        vendorParams: {
+            foo: "Bar"
+        },
+        vendorParamsShape: VendorParams
+    }
+])
+operation GetFoo {}
+
+structure VendorParams {
+    foo: String
+}