Allow mixins on all shape types

smithy-jsonschema/src/test/resources/software/amazon/smithy/jsonschema/model-with-mixins.smithy

@@ -7,6 +7,6 @@ structure Mixin {
     foo: String
 }
 
-structure UsesMixin with Mixin {
+structure UsesMixin with [Mixin] {
     baz: String
 }

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

@@ -228,7 +228,7 @@ private PendingShape createPendingShape(
             for (ShapeId mixin : mixins) {
                 Shape mixinShape = shapeMap.get(mixin);
                 for (MemberShape member : mixinShape.members()) {
-                    if (builderMembers.containsKey(member.getMemberName())) {
+                    if (memberConflicts(builderMembers, member)) {
                         // Members cannot be redefined.
                         MemberShape.Builder conflict = builderMembers.get(member.getMemberName());
                         events.add(ValidationEvent.builder()
@@ -260,6 +260,15 @@ private PendingShape createPendingShape(
         });
     }
 
+    private boolean memberConflicts(Map<String, MemberShape.Builder> builderMembers, MemberShape mixinMember) {
+        if (!builderMembers.containsKey(mixinMember.getMemberName())) {
+            return false;
+        }
+
+        MemberShape localMember = builderMembers.get(mixinMember.getMemberName()).build();
+        return !localMember.getTarget().equals(mixinMember.getTarget());
+    }
+
     private <S extends Shape, B extends AbstractShapeBuilder<? extends B, S>> Optional<S> buildShape(
             B builder,
             TraitContainer resolvedTraits

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

@@ -223,6 +223,10 @@ private void loadMember(FullyResolvedModelFile modelFile, ShapeId id, ObjectNode
         modelFile.onShape(builder);
     }
 
+    private void loadOptionalMember(FullyResolvedModelFile modelFile, ShapeId id, ObjectNode node, String member) {
+        node.getObjectMember(member).ifPresent(targetNode -> loadMember(modelFile, id, targetNode));
+    }
+
     private void loadCollection(
             ShapeId id,
             ObjectNode node,
@@ -231,16 +235,18 @@ private void loadCollection(
     ) {
         LoaderUtils.checkForAdditionalProperties(node, id, COLLECTION_PROPERTY_NAMES, modelFile.events());
         applyShapeTraits(id, node, modelFile);
-        loadMember(modelFile, id.withMember("member"), node.expectObjectMember("member"));
+        loadOptionalMember(modelFile, id.withMember("member"), node, "member");
         modelFile.onShape(builder.id(id).source(node.getSourceLocation()));
+        addMixins(id, node, modelFile);
     }
 
     private void loadMap(ShapeId id, ObjectNode node, FullyResolvedModelFile modelFile) {
         LoaderUtils.checkForAdditionalProperties(node, id, MAP_PROPERTY_NAMES, modelFile.events());
-        loadMember(modelFile, id.withMember("key"), node.expectObjectMember("key"));
-        loadMember(modelFile, id.withMember("value"), node.expectObjectMember("value"));
+        loadOptionalMember(modelFile, id.withMember("key"), node, "key");
+        loadOptionalMember(modelFile, id.withMember("value"), node, "value");
         applyShapeTraits(id, node, modelFile);
         modelFile.onShape(MapShape.builder().id(id).source(node.getSourceLocation()));
+        addMixins(id, node, modelFile);
     }
 
     private void loadOperation(ShapeId id, ObjectNode node, FullyResolvedModelFile modelFile) {
@@ -253,6 +259,7 @@ private void loadOperation(ShapeId id, ObjectNode node, FullyResolvedModelFile m
         loadOptionalTarget(modelFile, id, node, "input").ifPresent(builder::input);
         loadOptionalTarget(modelFile, id, node, "output").ifPresent(builder::output);
         modelFile.onShape(builder);
+        addMixins(id, node, modelFile);
     }
 
     private void loadResource(ShapeId id, ObjectNode node, FullyResolvedModelFile modelFile) {
@@ -279,6 +286,7 @@ private void loadResource(ShapeId id, ObjectNode node, FullyResolvedModelFile mo
         });
 
         modelFile.onShape(builder);
+        addMixins(id, node, modelFile);
     }
 
     private void loadService(ShapeId id, ObjectNode node, FullyResolvedModelFile modelFile) {
@@ -291,6 +299,7 @@ private void loadService(ShapeId id, ObjectNode node, FullyResolvedModelFile mod
         loadServiceRenameIntoBuilder(builder, node);
         builder.addErrors(loadOptionalTargetList(modelFile, id, node, ERRORS));
         modelFile.onShape(builder);
+        addMixins(id, node, modelFile);
     }
 
     static void loadServiceRenameIntoBuilder(ServiceShape.Builder builder, ObjectNode node) {
@@ -308,6 +317,7 @@ private void loadSimpleShape(
         LoaderUtils.checkForAdditionalProperties(node, id, SIMPLE_PROPERTY_NAMES, modelFile.events());
         applyShapeTraits(id, node, modelFile);
         modelFile.onShape(builder.id(id).source(node.getSourceLocation()));
+        addMixins(id, node, modelFile);
     }
 
     private void loadStructure(ShapeId id, ObjectNode node, FullyResolvedModelFile modelFile) {
@@ -328,7 +338,10 @@ private void finishLoadingStructOrUnionMembers(ShapeId id, ObjectNode node, Full
         for (Map.Entry<String, Node> entry : memberObject.getStringMap().entrySet()) {
             loadMember(modelFile, id.withMember(entry.getKey()), entry.getValue().expectObjectNode());
         }
+        addMixins(id, node, modelFile);
+    }
 
+    private void addMixins(ShapeId id, ObjectNode node, FullyResolvedModelFile modelFile) {
         ArrayNode mixins = node.getArrayMember(MIXINS).orElse(Node.arrayNode());
         for (ObjectNode mixin : mixins.getElementsAs(ObjectNode.class)) {
             modelFile.addPendingMixin(id, loadReferenceBody(modelFile, id, mixin));

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

@@ -66,7 +66,6 @@
 import software.amazon.smithy.model.traits.TraitFactory;
 import software.amazon.smithy.model.validation.Severity;
 import software.amazon.smithy.model.validation.ValidationEvent;
-import software.amazon.smithy.model.validation.ValidationUtils;
 import software.amazon.smithy.model.validation.Validator;
 import software.amazon.smithy.utils.ListUtils;
 import software.amazon.smithy.utils.SetUtils;
@@ -443,7 +442,6 @@ private void parseShape(List<TraitEntry> traits) {
         }
 
         addTraits(id, traits);
-        clearPendingDocs();
         ws();
     }
 
@@ -454,22 +452,25 @@ private ShapeId parseShapeName() {
 
     private void parseSimpleShape(ShapeId id, SourceLocation location, AbstractShapeBuilder builder) {
         modelFile.onShape(builder.source(location).id(id));
+        parseMixins(id);
+        // Pending docs get cleared by parseMixins. It can't be done here because the
+        // mixin parser needs to consume whitespace that may include doc comments for
+        // the next shape in the model.
     }
 
     // See parseMap for information on why members are parsed before the
     // list/set is registered with the ModelFile.
     private void parseCollection(ShapeId id, SourceLocation location, CollectionShape.Builder builder) {
         ws();
         builder.id(id).source(location);
+        parseMixins(id);
         parseMembers(id, SetUtils.of("member"));
         modelFile.onShape(builder.id(id));
+        clearPendingDocs();
     }
 
     private void parseMembers(ShapeId id, Set<String> requiredMembers) {
         Set<String> definedMembers = new HashSet<>();
-        Set<String> remaining = requiredMembers.isEmpty()
-                ? requiredMembers
-                : new HashSet<>(requiredMembers);
 
         ws();
         expect('{');
@@ -480,7 +481,7 @@ private void parseMembers(ShapeId id, Set<String> requiredMembers) {
                 break;
             }
 
-            parseMember(id, requiredMembers, definedMembers, remaining);
+            parseMember(id, requiredMembers, definedMembers);
 
             // Clears out any previously captured documentation
             // comments that may have been found when parsing the member.
@@ -493,15 +494,14 @@ private void parseMembers(ShapeId id, Set<String> requiredMembers) {
             expect('}');
         }
 
-        if (!remaining.isEmpty()) {
-            throw syntax(id, "Missing required members of shape `" + id + "`: ["
-                         + ValidationUtils.tickedList(remaining) + ']');
-        }
-
         expect('}');
     }
 
-    private void parseMember(ShapeId parent, Set<String> required, Set<String> defined, Set<String> remaining) {
+    private void parseMember(
+            ShapeId parent,
+            Set<String> allowed,
+            Set<String> defined
+    ) {
         // Parse optional member traits.
         List<TraitEntry> memberTraits = parseDocsAndTraits();
         SourceLocation memberLocation = currentLocation();
@@ -513,10 +513,9 @@ private void parseMember(ShapeId parent, Set<String> required, Set<String> defin
         }
 
         defined.add(memberName);
-        remaining.remove(memberName);
 
         // Only enforce "allowedMembers" if it isn't empty.
-        if (!required.isEmpty() && !required.contains(memberName)) {
+        if (!allowed.isEmpty() && !allowed.contains(memberName)) {
             throw syntax(parent, "Unexpected member of " + parent + ": '" + memberName + '\'');
         }
 
@@ -545,8 +544,10 @@ private void parseMapStatement(ShapeId id, SourceLocation location) {
         // on a builder. This does not suffer the same error messages as
         // structures/unions because list/set/map have a fixed and required
         // set of members that must be provided.
+        parseMixins(id);
         parseMembers(id, SetUtils.of("key", "value"));
         modelFile.onShape(MapShape.builder().id(id).source(location));
+        clearPendingDocs();
     }
 
     private void parseStructuredShape(
@@ -566,9 +567,13 @@ private void parseStructuredShape(
         // Parse optional "with" statements to add mixins, but only if it's supported by the version.
         parseMixins(id);
         parseMembers(id, Collections.emptySet());
+        clearPendingDocs();
     }
 
     private void parseMixins(ShapeId id) {
+        // This is fine for now, but if we ever add any other keywords that start with
+        // 'w' then we'll need to peek farther.
+        ws();
         if (peek() != 'w') {
             return;
         }
@@ -584,14 +589,20 @@ private void parseMixins(ShapeId id) {
         }
 
         ws();
+        expect('[');
+        ws();
+
         do {
+            clearPendingDocs();
             String target = ParserUtils.parseShapeId(this);
             modelFile.addForwardReference(target, resolved -> modelFile.addPendingMixin(id, resolved));
             ws();
-        } while (peek() != '{');
+        } while (peek() != ']');
+        expect(']');
     }
 
     private void parseOperationStatement(ShapeId id, SourceLocation location) {
+        parseMixins(id);
         OperationShape.Builder builder = OperationShape.builder().id(id).source(location);
         parseProperties(id, propertyName -> {
             switch (propertyName) {
@@ -611,6 +622,7 @@ private void parseOperationStatement(ShapeId id, SourceLocation location) {
             }
         });
         modelFile.onShape(builder);
+        clearPendingDocs();
     }
 
     private void parseProperties(ShapeId id, Consumer<String> valueParser) {
@@ -690,6 +702,8 @@ private void parseIdList(Consumer<ShapeId> consumer) {
     }
 
     private void parseServiceStatement(ShapeId id, SourceLocation location) {
+        ws();
+        parseMixins(id);
         ws();
         ServiceShape.Builder builder = new ServiceShape.Builder().id(id).source(location);
         ObjectNode shapeNode = IdlNodeParser.parseObjectNode(this, id.toString());
@@ -700,6 +714,7 @@ private void parseServiceStatement(ShapeId id, SourceLocation location) {
         optionalIdList(shapeNode, RESOURCES_KEY, builder::addResource);
         optionalIdList(shapeNode, ERRORS_KEYS, builder::addError);
         AstModelLoader.loadServiceRenameIntoBuilder(builder, shapeNode);
+        clearPendingDocs();
     }
 
     private void optionalId(ObjectNode node, String name, Consumer<ShapeId> consumer) {
@@ -718,6 +733,8 @@ private void optionalIdList(ObjectNode node, String name, Consumer<ShapeId> cons
     }
 
     private void parseResourceStatement(ShapeId id, SourceLocation location) {
+        ws();
+        parseMixins(id);
         ws();
         ResourceShape.Builder builder = ResourceShape.builder().id(id).source(location);
         modelFile.onShape(builder);
@@ -742,6 +759,7 @@ private void parseResourceStatement(ShapeId id, SourceLocation location) {
                 modelFile.addForwardReference(target.getValue(), targetId -> builder.addIdentifier(name, targetId));
             }
         });
+        clearPendingDocs();
     }
 
     // "//" *(not_newline)

smithy-model/src/main/java/software/amazon/smithy/model/shapes/CollectionShape.java

@@ -19,7 +19,8 @@
 import java.util.Collections;
 import java.util.Objects;
 import java.util.function.Consumer;
-import software.amazon.smithy.utils.SmithyBuilder;
+import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.traits.Trait;
 
 /**
  * Abstract class representing Set and List shapes.
@@ -28,9 +29,9 @@ public abstract class CollectionShape extends Shape {
 
     private final MemberShape member;
 
-    CollectionShape(Builder builder) {
+    CollectionShape(Builder<?, ?> builder) {
         super(builder, false);
-        member = SmithyBuilder.requiredState("member", builder.member);
+        member = builder.member != null ? builder.member : getRequiredMixinMember(builder, "member");
         ShapeId expected = getId().withMember("member");
         if (!member.getId().equals(expected)) {
             throw new IllegalArgumentException(String.format(
@@ -124,5 +125,27 @@ public B member(ShapeId target, Consumer<MemberShape.Builder> memberUpdater) {
         public final B addMember(MemberShape member) {
             return member(member);
         }
+
+        @Override
+        public B flattenMixins() {
+            for (Shape mixin : getMixins().values()) {
+                SourceLocation location = getSourceLocation();
+                Collection<Trait> localTraits = Collections.emptyList();
+                MemberShape mixinMember = ((CollectionShape) mixin).getMember();
+                MemberShape existing = member;
+                if (existing != null) {
+                    localTraits = existing.getIntroducedTraits().values();
+                    location = existing.getSourceLocation();
+                }
+                member = MemberShape.builder()
+                        .id(getId().withMember(mixinMember.getMemberName()))
+                        .target(mixinMember.getTarget())
+                        .addTraits(mixinMember.getAllTraits().values())
+                        .addTraits(localTraits)
+                        .source(location)
+                        .build();
+            }
+            return super.flattenMixins();
+        }
     }
 }