Repository: smithy-lang/smithy
Pull request #2378 (2024-08-23 15:56:18)

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 we apply the mixin.

---
By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.



Patch:
diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/loader/ApplyMixin.java b/smithy-model/src/main/java/software/amazon/smithy/model/loader/ApplyMixin.java
index 647f86b2b..2f7791c6a 100644
--- a/smithy-model/src/main/java/software/amazon/smithy/model/loader/ApplyMixin.java
+++ b/smithy-model/src/main/java/software/amazon/smithy/model/loader/ApplyMixin.java
@@ -17,6 +17,8 @@ package software.amazon.smithy.model.loader;
 
 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 @@ import software.amazon.smithy.model.validation.Validator;
 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 @@ final class ApplyMixin implements ShapeModifier {
 
         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 @@ final class ApplyMixin implements ShapeModifier {
     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);
+    }
 }
diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/loader/LoaderTraitMap.java b/smithy-model/src/main/java/software/amazon/smithy/model/loader/LoaderTraitMap.java
index 639cebc88..155f5f77b 100644
--- a/smithy-model/src/main/java/software/amazon/smithy/model/loader/LoaderTraitMap.java
+++ b/smithy-model/src/main/java/software/amazon/smithy/model/loader/LoaderTraitMap.java
@@ -83,10 +83,19 @@ final class LoaderTraitMap {
                     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<>())


Test patch:
diff --git a/smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java b/smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java
index 3c542b41f..ad4ce29eb 100644
--- a/smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java
+++ b/smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java
@@ -1416,4 +1416,21 @@ public class ModelAssemblerTest {
                 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));
+    }
 }
diff --git a/smithy-model/src/test/resources/software/amazon/smithy/model/loader/mixin-and-apply-model.smithy b/smithy-model/src/test/resources/software/amazon/smithy/model/loader/mixin-and-apply-model.smithy
new file mode 100644
index 000000000..228a58f3e
--- /dev/null
+++ b/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