Fix multi-model, mixin member apply bug

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

@@ -17,6 +17,8 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -37,6 +39,7 @@
 final class ApplyMixin implements ShapeModifier {
     private final ShapeId mixin;
     private List<ValidationEvent> events;
+    private Map<ShapeId, Map<ShapeId, Trait>> potentiallyIntroducedTraits;
 
     ApplyMixin(ShapeId mixin) {
         this.mixin = mixin;
@@ -80,8 +83,11 @@ public void modifyShape(
 
         for (MemberShape member : mixinShape.members()) {
             ShapeId targetId = builder.getId().withMember(member.getMemberName());
-            // Claim traits from the trait map that were applied to synthesized shapes.
-            Map<ShapeId, Trait> introducedTraits = unclaimedTraits.apply(targetId);
+            // Claim traits from the trait maps that were applied to synthesized shapes.
+            Map<ShapeId, Trait> introducedTraits = new LinkedHashMap<>(unclaimedTraits.apply(targetId));
+            if (potentiallyIntroducedTraits != null && potentiallyIntroducedTraits.containsKey(targetId)) {
+                introducedTraits.putAll(potentiallyIntroducedTraits.get(targetId));
+            }
             String memberName = member.getMemberName();
             MemberShape introducedMember = null;
             Optional<MemberShape> previouslyAdded = builder.getMember(memberName);
@@ -145,4 +151,14 @@ private void mixinMemberConflict(MemberShape.Builder conflict, MemberShape other
     public List<ValidationEvent> getEvents() {
         return events == null ? Collections.emptyList() : events;
     }
+
+    void putPotentiallyIntroducedTrait(ShapeId target, Trait trait) {
+        if (potentiallyIntroducedTraits == null) {
+            potentiallyIntroducedTraits = new HashMap<>();
+        }
+
+        Map<ShapeId, Trait> shapeUnclaimedTraits = potentiallyIntroducedTraits.computeIfAbsent(target,
+                id -> new LinkedHashMap<>());
+        shapeUnclaimedTraits.put(trait.toShapeId(), trait);
+    }
 }

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

@@ -83,10 +83,19 @@ void applyTraitsToNonMixinsInShapeMap(LoaderShapeMap shapeMap) {
                     for (LoadOperation.DefineShape shape : rootShapes) {
                         if (shape.hasMember(memberName)) {
                             foundMember = true;
-                            shape.memberBuilders().get(memberName).getAllTraits();
                             applyTraitsToShape(shape.memberBuilders().get(memberName), created);
+                        } else {
+                            // If we didn't have the member and the member is from a mixin,
+                            // we need to update ApplyMixin shape modifiers to apply the trait
+                            // in case we have the target's container already in the shapeMap.
+                            for (ShapeModifier modifier : shape.modifiers()) {
+                                if (modifier instanceof ApplyMixin) {
+                                    ((ApplyMixin) modifier).putPotentiallyIntroducedTrait(target, created);
+                                }
+                            }
                         }
                     }
+
                     // If the member wasn't found, then it might be a mixin member that is synthesized later.
                     if (!foundMember) {
                         unclaimed.computeIfAbsent(target.withMember(memberName), id -> new LinkedHashMap<>())

smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java

@@ -1416,4 +1416,21 @@ public void loadsShapesWhenThereAreUnresolvedMixins() {
                 ShapeId.from("com.foo#Bar")
         ));
     }
+
+    @Test
+    public void handlesSameModelWhenBuiltAndImported() throws Exception {
+        Path modelUri = Paths.get(getClass().getResource("mixin-and-apply-model.smithy").toURI());
+        Model sourceModel = Model.assembler()
+                .addImport(modelUri)
+                .assemble()
+                .unwrap();
+        Model combinedModel = Model.assembler()
+                .addModel(sourceModel)
+                .addImport(modelUri)
+                .assemble()
+                .unwrap();
+
+        assertTrue(combinedModel.expectShape(ShapeId.from("smithy.example#MachineData$machineId"), MemberShape.class)
+                .hasTrait(RequiredTrait.ID));
+    }
 }

smithy-model/src/test/resources/software/amazon/smithy/model/loader/mixin-and-apply-model.smithy

@@ -0,0 +1,23 @@
+$version: "2.0"
+$operationInputSuffix: "Request"
+$operationOutputSuffix: "Response"
+
+namespace smithy.example
+
+@http(uri: "/machines", method: "POST")
+operation CreateMachine {
+    output := {
+        Machine: MachineData
+    }
+}
+
+@mixin
+structure MachineDataMixin {
+    machineId: String
+}
+
+structure MachineData with [MachineDataMixin] {
+    manufacturer: String
+}
+
+apply MachineData$machineId @required