Fix multi-model, mixin member apply bug

On first load, via sources, there's a `DefineShape` `LoadOperation`
with an `ApplyMixin` modifier. There's also a pending `@apply`
statement in the `LoaderTraitMap`. When the source is assembled,
that's fully resolved, the mixin is merged into the shape, then
traits applied to mixed content are applied. This is a valid source
model that has the trait applied to the mixed member.

The import is then broken out into `LoadOperation`s to be assembled.
The shape's `LoadOperations` already exist in the map and a new
`DefineShape` load operation is added that also has the `ApplyMixin`
modifier. The first, already fully mixed applied and built - the
second, unmixed and unapplied.

When we go to `applyTraitsToNonMixinsInShapeMap` on the import pass,
we detect if any of the load operations that `DefineShape` contain
the member where the trait is applied and add it. However, because
we have **two** `LoadOperations` - one with the member and one
without - we apply the trait to the shape that was already mixed and
applied. We do not add it to the second `LoadOperation`, because it
does not have the member. But, because we applied it somewhere, the
trait is no longer unclaimed. When then later resolving the
`ApplyMixin` modifier, the trait is not applied since it has been
claimed elsewhere. When building the second shape and comparing it
to the first, the trait is then missing, failing out due to a conflict.

This is resolved by sending traits for missing members directly in to the
`ApplyMixin` modifier and merging them when apply the mixin.

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