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 647f86b2b01..2f7791c6a05 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 @@ 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 events; + private Map> 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 introducedTraits = unclaimedTraits.apply(targetId); + // Claim traits from the trait maps that were applied to synthesized shapes. + Map 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 previouslyAdded = builder.getMember(memberName); @@ -145,4 +151,14 @@ private void mixinMemberConflict(MemberShape.Builder conflict, MemberShape other public List getEvents() { return events == null ? Collections.emptyList() : events; } + + void putPotentiallyIntroducedTrait(ShapeId target, Trait trait) { + if (potentiallyIntroducedTraits == null) { + potentiallyIntroducedTraits = new HashMap<>(); + } + + Map 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 639cebc889f..155f5f77b90 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 @@ 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<>()) 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 3c542b41f9d..ad4ce29ebd4 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 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)); + } } 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 00000000000..228a58f3e9e --- /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