Fix replacing shapes when traits applied to mixed-in members (#1509)

smithy-build/src/test/java/software/amazon/smithy/build/transforms/ExcludeShapesByTagTest.java

@@ -18,6 +18,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 
 import java.nio.file.Paths;
 import java.util.Optional;
@@ -27,6 +28,7 @@
 import software.amazon.smithy.model.node.Node;
 import software.amazon.smithy.model.shapes.EnumShape;
 import software.amazon.smithy.model.shapes.IntEnumShape;
+import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.StringShape;
 import software.amazon.smithy.model.traits.TagsTrait;
@@ -73,5 +75,9 @@ public void filtersMembers() throws Exception {
 
         IntEnumShape bar = result.expectShape(ShapeId.from("smithy.example#Bar"), IntEnumShape.class);
         assertThat(bar.members().size(), is(1));
+
+        // Mixin members are retained, but excluded traits are removed.
+        MemberShape baz = result.expectShape(ShapeId.from("smithy.example#StructForMixin$baz"), MemberShape.class);
+        assertFalse(baz.findMemberTrait(result, "MyTrait").isPresent());
     }
 }

smithy-build/src/test/java/software/amazon/smithy/build/transforms/ExcludeTraitsByTagTest.java

@@ -25,6 +25,7 @@
 import software.amazon.smithy.build.TransformContext;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.TraitDefinition;
@@ -47,5 +48,11 @@ public void removesTraitsByTagInList() throws Exception {
 
         assertFalse(traits.contains(ShapeId.from("ns.foo#quux")));
         assertTrue(traits.contains(ShapeId.from("ns.foo#bar")));
+
+        // Mixin members are retained, but tagged traits are excluded.
+        MemberShape mixedMember = result.expectShape(ShapeId.from("ns.foo#MyOperationInput$mixedMember"),
+                MemberShape.class);
+        assertFalse(mixedMember.findMemberTrait(result, "ns.foo#corge").isPresent());
+        assertTrue(mixedMember.findMemberTrait(result, "ns.foo#bar").isPresent());
     }
 }

smithy-build/src/test/resources/software/amazon/smithy/build/transforms/filter-by-tags.smithy

@@ -19,3 +19,19 @@ intEnum Bar {
     @enumValue(2)
     KEEP
 }
+
+@trait
+@tags(["filter"])
+string MyTrait
+
+@mixin
+structure MyMixin {
+    @required
+    baz: String
+}
+
+structure StructForMixin with [MyMixin] {
+    bar: String
+}
+
+apply StructForMixin$baz @MyTrait("hello")

smithy-build/src/test/resources/software/amazon/smithy/build/transforms/tree-shaking-traits.json

@@ -31,6 +31,21 @@
                 ]
             }
         },
+        "ns.foo#corge": {
+            "type": "structure",
+            "members": {
+                "member": {
+                    "target": "ns.foo#CorgeTraitShapeMember"
+                }
+            },
+            "traits": {
+                "smithy.api#trait": {},
+                "smithy.api#tags": [
+                    "foo",
+                    "qux"
+                ]
+            }
+        },
         "ns.foo#MyService": {
             "type": "service",
             "version": "2017-01-19",
@@ -46,6 +61,22 @@
                 "target": "ns.foo#MyOperationInput"
             }
         },
+        "ns.foo#MyMixin": {
+            "type": "structure",
+            "members": {
+                "mixedMember": {
+                    "target": "smithy.api#String",
+                    "traits": {
+                        "ns.foo#bar": {
+                            "member": "baz"
+                        }
+                    }
+                }
+            },
+            "traits": {
+                "smithy.api#mixin": {}
+            }
+        },
         "ns.foo#MyOperationInput": {
             "type": "structure",
             "members": {
@@ -55,7 +86,12 @@
                 "buzz": {
                     "target": "ns.foo#Include2"
                 }
-            }
+            },
+            "mixins": [
+                {
+                    "target": "ns.foo#MyMixin"
+                }
+            ]
         },
         "ns.foo#Exclude1": {
             "type": "string",
@@ -84,6 +120,17 @@
         },
         "ns.foo#QuuxTraitShapeMember": {
             "type": "string"
+        },
+        "ns.foo#CorgeTraitShapeMember": {
+            "type": "string"
+        },
+        "ns.foo#MyOperationInput$mixedMember": {
+            "type": "apply",
+            "traits": {
+                "ns.foo#corge": {
+                    "member": "hi"
+                }
+            }
         }
     }
 }

smithy-model/src/main/java/software/amazon/smithy/model/loader/TopologicalShapeSort.java

@@ -32,7 +32,7 @@
 /**
  * Topologically sorts shapes based on their dependencies (i.e., mixins).
  *
- * <p>While this class is reusable, is is also stateful; shapes and edges
+ * <p>While this class is reusable, it is also stateful; shapes and edges
  * are enqueued, and when sorted, all shapes and edges are dequeued.
  */
 public final class TopologicalShapeSort {

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

@@ -20,6 +20,7 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -189,7 +190,12 @@ private void updateMixins(Model model, Model.Builder builder, List<Shape> replac
 
         // Add _all_ of the replacements in case mixins or the Mixin trait were removed from updated shapes.
         for (Shape shape : replacements) {
-            sorter.enqueue(shape);
+            // Enqueue member shapes with empty dependencies since the parent will be enqueued with them.
+            if (shape.isMemberShape()) {
+                sorter.enqueue(shape.getId(), Collections.emptySet());
+            } else {
+                sorter.enqueue(shape);
+            }
         }
 
         List<ShapeId> sorted = sorter.dequeueSortedShapes();