Move IDL specific loading logic to IDL

This commit moves all IDL-specific model loading logic into the IDL
loader and out of the LoaderVisitor. This cleans up the LoaderVisitor
and also fixes a bug in how the `apply` statement wasn't properly
resolving the targeted shape ID (it always assumed that the target was
in the current namespace).

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

@@ -87,7 +87,6 @@ enum AstModelLoader {
 
     void load(ObjectNode model, LoaderVisitor visitor) {
         model.expectNoAdditionalProperties(TOP_LEVEL_PROPERTIES);
-        visitor.onOpenFile(model.getSourceLocation().getFilename());
         loadMetadata(model, visitor);
         loadShapes(model, visitor);
     }

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

@@ -42,6 +42,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -52,6 +53,8 @@
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
+import software.amazon.smithy.model.FromSourceLocation;
+import software.amazon.smithy.model.SourceException;
 import software.amazon.smithy.model.SourceLocation;
 import software.amazon.smithy.model.node.ArrayNode;
 import software.amazon.smithy.model.node.BooleanNode;
@@ -141,6 +144,9 @@ final class IdlModelLoader {
     private final LoaderVisitor visitor;
     private final List<Pair<String, Node>> pendingTraits = new ArrayList<>();
 
+    /** Map of shape aliases to their targets. */
+    private final Map<String, ShapeId> useShapes = new HashMap<>();
+
     private Token current;
     private DocComment pendingDocComment;
     private String definedVersion;
@@ -152,7 +158,6 @@ private IdlModelLoader(String filename, SmithyModelLexer lexer, LoaderVisitor vi
         this.filename = filename;
         this.visitor = visitor;
         this.lexer = lexer;
-        visitor.onOpenFile(filename);
 
         while (!eof()) {
             Token token = expect(UNQUOTED, ANNOTATION, CONTROL, DOC);
@@ -285,8 +290,7 @@ private void parseNamespace() {
             throw syntax(format("Invalid namespace name `%s`", parsedNamespace));
         }
 
-        visitor.onNamespace(parsedNamespace, current());
-        this.namespace = parsedNamespace;
+        namespace = parsedNamespace;
     }
 
     private void parseUseStatement() {
@@ -300,7 +304,13 @@ private void parseUseStatement() {
 
         try {
             Token namespaceToken = expect(UNQUOTED);
-            visitor.onUseShape(ShapeId.from(namespaceToken.lexeme), namespaceToken);
+            ShapeId target = ShapeId.from(namespaceToken.lexeme);
+            ShapeId previous = useShapes.put(target.getName(), target);
+            if (previous != null) {
+                String message = String.format("Cannot use name `%s` because it conflicts with `%s`",
+                                               target, previous);
+                throw new UseException(message, namespaceToken.getSourceLocation());
+            }
             expectNewline();
         } catch (ShapeIdSyntaxException e) {
             throw syntax(e.getMessage());
@@ -385,7 +395,7 @@ private Pair<String, Node> parseTraitValue(Token token, TraitValueType type) {
             requireNamespaceOrThrow();
 
             // Resolve the trait name and ensure that the trait forms a syntactically valid value.
-            ShapeId.fromOptionalNamespace(visitor.getNamespace(), token.lexeme);
+            ShapeId.fromOptionalNamespace(namespace, token.lexeme);
             Pair<String, Node> result = Pair.of(token.lexeme, parseTraitValueBody());
 
             // `apply` doesn't require any specific token to follow.
@@ -481,11 +491,59 @@ private Node parseUnquotedNode(Token token) {
                 Pair<StringNode, Consumer<String>> pair = StringNode.createLazyString(
                         token.lexeme, token.getSourceLocation());
                 Consumer<String> consumer = pair.right;
-                visitor.onShapeTarget(token.lexeme, token, id -> consumer.accept(id.toString()));
+                onShapeTarget(token.lexeme, token, id -> consumer.accept(id.toString()));
                 return pair.left;
         }
     }
 
+    /**
+     * Resolve shape targets and tracks forward references.
+     *
+     * <p>Smithy models allow for forward references to shapes that have not
+     * yet been defined. Only after all shapes are loaded is the entire set
+     * of possible shape IDs known. This normally isn't a concern, but Smithy
+     * allows for public shapes defined in the prelude to be referenced by
+     * targets like members and resource identifiers without an absolute
+     * shape ID (for example, {@code String}). However, a relative prelude
+     * shape ID is only resolved for such a target if a shape of the same
+     * name was not defined in the same namespace in which the target
+     * was defined.
+     *
+     * <p>If a shape in the same namespace as the target has already been
+     * defined or if the target is absolute and cannot resolve to a prelude
+     * shape, the provided {@code resolver} is invoked immediately. Otherwise,
+     * the {@code resolver} is invoked inside of the {@link LoaderVisitor#onEnd}
+     * method only after all shapes have been declared.
+     *
+     * @param target Shape that is targeted.
+     * @param sourceLocation The location of where the target occurred.
+     * @param resolver The consumer to invoke once the shape ID is resolved.
+     */
+    private void onShapeTarget(String target, FromSourceLocation sourceLocation, Consumer<ShapeId> resolver) {
+        // Account for aliased shapes.
+        if (useShapes.containsKey(target)) {
+            resolver.accept(useShapes.get(target));
+            return;
+        }
+
+        try {
+            // A namespace is not set when parsing metadata.
+            ShapeId expectedId = namespace == null
+                                 ? ShapeId.from(target)
+                                 : ShapeId.fromOptionalNamespace(namespace, target);
+
+            if (isRealizedShapeId(expectedId, target)) {
+                // Account for previously seen shapes in this namespace, absolute shapes, and prelude namespaces
+                // always resolve to prelude shapes.
+                resolver.accept(expectedId);
+            } else {
+                visitor.addForwardReference(expectedId, resolver);
+            }
+        } catch (ShapeIdSyntaxException e) {
+            throw new SourceException("Error resolving shape target; " + e.getMessage(), sourceLocation, e);
+        }
+    }
+
     private NumberNode parseNumber(Token token) {
         if (token.lexeme.contains("e") || token.lexeme.contains(".")) {
             return new NumberNode(Double.valueOf(token.lexeme), token.getSourceLocation());
@@ -546,14 +604,49 @@ private void parseSimpleShape(AbstractShapeBuilder builder) {
      * @return Returns the parsed shape ID.
      */
     private ShapeId parseShapeName() {
+        requireNamespaceOrThrow();
         definedShapes = true;
         Token nameToken = expect(UNQUOTED);
         String name = nameToken.lexeme;
-        ShapeId id = visitor.onShapeDefName(name, nameToken);
-        pendingTraits.forEach(pair -> visitor.onTrait(id, pair.getLeft(), pair.getRight()));
-        pendingTraits.clear();
-        collectPendingDocString(id);
-        return id;
+
+        if (useShapes.containsKey(name)) {
+            String msg = String.format("shape name `%s` conflicts with imported shape `%s`",
+                                       name, useShapes.get(name));
+            throw new UseException(msg, nameToken);
+        }
+
+        try {
+            ShapeId id = ShapeId.fromRelative(namespace, name);
+            for (Pair<String, Node> pair : pendingTraits) {
+                onDeferredTrait(id, pair.getLeft(), pair.getRight());
+            }
+            pendingTraits.clear();
+            collectPendingDocString(id);
+            return id;
+        } catch (ShapeIdSyntaxException e) {
+            throw new ModelSyntaxException("Invalid shape name: " + name, nameToken);
+        }
+    }
+
+    /**
+     * Adds a trait to a shape after resolving all shape IDs.
+     *
+     * <p>Resolving the trait against a trait definition is deferred until
+     * the entire model is loaded. A namespace is required to have been set
+     * if the trait name is not absolute.
+     *
+     * @param target Shape to add the trait to.
+     * @param traitName Trait name to add.
+     * @param traitValue Trait value as a Node object.
+     */
+    private void onDeferredTrait(ShapeId target, String traitName, Node traitValue) {
+        onShapeTarget(traitName, traitValue.getSourceLocation(), id -> visitor.onTrait(target, id, traitValue));
+    }
+
+    private boolean isRealizedShapeId(ShapeId expectedId, String target) {
+        return Objects.equals(namespace, Prelude.NAMESPACE)
+               || visitor.hasDefinedShape(expectedId)
+               || target.contains("#");
     }
 
     private void collectPendingDocString(ShapeId id) {
@@ -606,7 +699,7 @@ private void parseStructuredContents(String shapeType, ShapeId parent, Collectio
                 parseMember(memberId);
                 // Add the loaded traits on the member now that the ID is known.
                 for (Pair<String, Node> pair : memberTraits) {
-                    visitor.onTrait(memberId, pair.getLeft(), pair.getRight());
+                    onDeferredTrait(memberId, pair.getLeft(), pair.getRight());
                 }
                 memberTraits.clear();
                 collectPendingDocString(memberId);
@@ -630,7 +723,7 @@ private void parseMember(ShapeId memberId) {
         MemberShape.Builder memberBuilder = MemberShape.builder().id(memberId).source(currentLocation());
         visitor.onShape(memberBuilder);
         Token targetToken = expect(UNQUOTED, QUOTED);
-        visitor.onShapeTarget(targetToken.lexeme, targetToken, memberBuilder::target);
+        onShapeTarget(targetToken.lexeme, targetToken, memberBuilder::target);
     }
 
     private void parseCollection(String shapeType, CollectionShape.Builder builder) {
@@ -658,12 +751,17 @@ private void parseMap() {
     private void parseApply() {
         requireNamespaceOrThrow();
         // apply <ShapeName> @<trait>\n
-        String name = expect(UNQUOTED).lexeme;
-        ShapeId id = ShapeId.fromOptionalNamespace(visitor.getNamespace(), name);
+        Token nextToken = expect(UNQUOTED);
+        String name = nextToken.lexeme;
         Token token = expect(ANNOTATION);
         Pair<String, Node> trait = parseTraitValue(token, TraitValueType.APPLY);
-        visitor.onTrait(id, trait.getLeft(), trait.getRight());
         expectNewline();
+
+        // First, resolve the targeted shape.
+        onShapeTarget(name, nextToken.getSourceLocation(), id -> {
+            // Next, resolve the trait ID.
+            onDeferredTrait(id, trait.getLeft(), trait.getRight());
+        });
     }
 
     /**
@@ -833,7 +931,7 @@ private void parseResource() {
             for (Map.Entry<StringNode, Node> entry : ids.getMembers().entrySet()) {
                 String name = entry.getKey().getValue();
                 StringNode target = entry.getValue().expectStringNode();
-                visitor.onShapeTarget(target.getValue(), target, id -> builder.addIdentifier(name, id));
+                onShapeTarget(target.getValue(), target, id -> builder.addIdentifier(name, id));
             }
         });
 
@@ -853,14 +951,14 @@ private void parseOperation() {
             ObjectNode node = parseObjectNode(opening.getSourceLocation(), RBRACE);
             node.expectNoAdditionalProperties(OPERATION_PROPERTY_NAMES);
             node.getStringMember("input").ifPresent(input -> {
-                visitor.onShapeTarget(input.getValue(), input, builder::input);
+                onShapeTarget(input.getValue(), input, builder::input);
             });
             node.getStringMember("output").ifPresent(output -> {
-                visitor.onShapeTarget(output.getValue(), output, builder::output);
+                onShapeTarget(output.getValue(), output, builder::output);
             });
             node.getArrayMember("errors").ifPresent(errors -> {
                 for (StringNode value : errors.getElementsAs(StringNode.class)) {
-                    visitor.onShapeTarget(value.getValue(), value, builder::addError);
+                    onShapeTarget(value.getValue(), value, builder::addError);
                 }
             });
         }
@@ -879,15 +977,15 @@ private void parseDeprecatedOperationSyntax(OperationShape.Builder builder) {
 
         Token next = expect(RPAREN, UNQUOTED);
         if (next.type == UNQUOTED) {
-            visitor.onShapeTarget(next.lexeme, next, builder::input);
+            onShapeTarget(next.lexeme, next, builder::input);
             expect(RPAREN);
         }
 
         // Parse the optionally present return value.
         if (test(RETURN)) {
             expect(RETURN);
             Token returnToken = expect(UNQUOTED);
-            visitor.onShapeTarget(returnToken.lexeme, returnToken, builder::output);
+            onShapeTarget(returnToken.lexeme, returnToken, builder::output);
         }
 
         // Parse the optionally present errors list.
@@ -897,7 +995,7 @@ private void parseDeprecatedOperationSyntax(OperationShape.Builder builder) {
             if (!test(RBRACKET)) {
                 while (true) {
                     Token errorToken = expect(UNQUOTED);
-                    visitor.onShapeTarget(errorToken.lexeme, errorToken, builder::addError);
+                    onShapeTarget(errorToken.lexeme, errorToken, builder::addError);
                     if (test(RBRACKET)) {
                         break;
                     }

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

@@ -36,7 +36,6 @@
 import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
-import software.amazon.smithy.model.shapes.ShapeIdSyntaxException;
 import software.amazon.smithy.model.shapes.ShapeType;
 import software.amazon.smithy.model.traits.DynamicTrait;
 import software.amazon.smithy.model.traits.Trait;
@@ -87,12 +86,6 @@ final class LoaderVisitor {
     /** Built trait definitions. */
     private final Map<ShapeId, TraitDefinition> builtTraitDefinitions = new HashMap<>();
 
-    /** Current namespace being parsed. */
-    private String namespace;
-
-    /** Map of shape aliases to their alias. */
-    private final Map<String, ShapeId> useShapes = new HashMap<>();
-
     /** The result that is populated when onEnd is called. */
     private ValidatedResult<Model> result;
 
@@ -152,97 +145,6 @@ public boolean hasDefinedShape(ShapeId id) {
         return builtShapes.containsKey(id) || pendingShapes.containsKey(id);
     }
 
-    /**
-     * Invoked each time a new file is loaded, causing any file-specific
-     * caches and metadata to be flushed.
-     *
-     * <p>More specifically, this method clears out the previously used
-     * trait and shape name aliases so that they do not affect the next
-     * file that is loaded.
-     *
-     * @param filename The name of the file being opened.
-     */
-    public void onOpenFile(String filename) {
-        LOGGER.fine(() -> "Beginning to parse " + filename);
-        namespace = null;
-        useShapes.clear();
-    }
-
-    /**
-     * Gets the namespace that is currently being loaded.
-     *
-     * @return Returns the namespace or null if none is set.
-     */
-    public String getNamespace() {
-        return namespace;
-    }
-
-    /**
-     * Sets and validates the current namespace of the visitor.
-     *
-     * <p>Relative shape and trait definitions will use this
-     * namespace by default.
-     *
-     * @param namespace Namespace being entered.
-     * @param source The source location of where the namespace is defined.
-     */
-    public void onNamespace(String namespace, FromSourceLocation source) {
-        if (!ShapeId.isValidNamespace(namespace)) {
-            String msg = String.format("Invalid namespace name `%s`", namespace);
-            throw new ModelSyntaxException(msg, source);
-        }
-
-        this.namespace = namespace;
-    }
-
-    /**
-     * Resolve shape targets and tracks forward references.
-     *
-     * <p>Smithy models allow for forward references to shapes that have not
-     * yet been defined. Only after all shapes are loaded is the entire set
-     * of possible shape IDs known. This normally isn't a concern, but Smithy
-     * allows for public shapes defined in the prelude to be referenced by
-     * targets like members and resource identifiers without an absolute
-     * shape ID (for example, {@code String}). However, a relative prelude
-     * shape ID is only resolved for such a target if a shape of the same
-     * name was not defined in the same namespace in which the target
-     * was defined.
-     *
-     * <p>If a shape in the same namespace as the target has already been
-     * defined or if the target is absolute and cannot resolve to a prelude
-     * shape, the provided {@code resolver} is invoked immediately. Otherwise,
-     * the {@code resolver} is invoked inside of the {@link #onEnd} method
-     * only after all shapes have been declared.
-     *
-     * @param target Shape that is targeted.
-     * @param sourceLocation The location of where the target occurred.
-     * @param resolver The consumer to invoke once the shape ID is resolved.
-     */
-    public void onShapeTarget(String target, FromSourceLocation sourceLocation, Consumer<ShapeId> resolver) {
-        try {
-            // Account for aliased shapes.
-            if (useShapes.containsKey(target)) {
-                resolver.accept(useShapes.get(target));
-                return;
-            }
-
-            // A namespace is not set when parsing metadata.
-            ShapeId expectedId = namespace == null
-                    ? ShapeId.from(target)
-                    : ShapeId.fromOptionalNamespace(namespace, target);
-
-            if (Objects.equals(namespace, Prelude.NAMESPACE) || hasDefinedShape(expectedId) || target.contains("#")) {
-                // Account for previously seen shapes in this namespace, absolute shapes, and prelude namespaces
-                // always resolve to prelude shapes.
-                resolver.accept(expectedId);
-            } else {
-                forwardReferenceResolvers.add(new ForwardReferenceResolver(expectedId, resolver));
-            }
-        } catch (ShapeIdSyntaxException e) {
-            throw new SourceException("Error resolving shape target; " + e.getMessage(), sourceLocation, e);
-        }
-    }
-
     /**
      * Checks if a specific property is set.
      *
@@ -284,12 +186,6 @@ public <T> Optional<T> getProperty(String property, Class<T> type) {
         });
     }
 
-    private void validateState(FromSourceLocation sourceLocation) {
-        if (result != null) {
-            throw new IllegalStateException("Cannot call visitor method because visitor has called onEnd");
-        }
-    }
-
     /**
      * Adds an error to the loader.
      *
@@ -299,50 +195,12 @@ public void onError(ValidationEvent event) {
         events.add(Objects.requireNonNull(event));
     }
 
-    /**
-     * Invoked when a shape definition name is defined, validates the name,
-     * and returns the resolved shape ID of the name.
-     *
-     * <p>This method ensures a namespace has been set, the syntax is valid,
-     * and that the name does not conflict with any previously defined use
-     * statements.
-     *
-     * <p>This method has no side-effects.
-     *
-     * @param name Name being defined.
-     * @param source The location of where it is defined.
-     * @return Returns the parsed and loaded shape ID.
-     */
-    public ShapeId onShapeDefName(String name, FromSourceLocation source) {
-        validateState(source);
-        assertNamespaceIsPresent(source);
-
-        if (useShapes.containsKey(name)) {
-            String msg = String.format("shape name `%s` conflicts with imported shape `%s`",
-                                       name, useShapes.get(name));
-            throw new UseException(msg, source);
-        }
-
-        try {
-            return ShapeId.fromRelative(namespace, name);
-        } catch (ShapeIdSyntaxException e) {
-            throw new ModelSyntaxException("Invalid shape name: " + name, source);
-        }
-    }
-
-    private void assertNamespaceIsPresent(FromSourceLocation source) {
-        if (namespace == null) {
-            throw new ModelSyntaxException("A namespace must be set before shapes or traits can be defined", source);
-        }
-    }
-
     /**
      * Adds a shape to the loader.
      *
      * @param shapeBuilder Shape builder to add.
      */
     public void onShape(AbstractShapeBuilder shapeBuilder) {
-        validateState(shapeBuilder);
         ShapeId id = SmithyBuilder.requiredState("id", shapeBuilder.getId());
         if (validateOnShape(id, shapeBuilder)) {
             pendingShapes.put(id, shapeBuilder);
@@ -355,7 +213,6 @@ public void onShape(AbstractShapeBuilder shapeBuilder) {
      * @param shape Built shape to add to the loader visitor.
      */
     public void onShape(Shape shape) {
-        validateState(shape);
         if (validateOnShape(shape.getId(), shape)) {
             builtShapes.put(shape.getId(), shape);
         }
@@ -403,21 +260,6 @@ private boolean validateOnShape(ShapeId id, FromSourceLocation source) {
         return false;
     }
 
-    /**
-     * Adds a trait to a shape.
-     *
-     * <p>Resolving the trait against a trait definition is deferred until
-     * the entire model is loaded. A namespace is required to have been set
-     * if the trait name is not absolute.
-     *
-     * @param target Shape to add the trait to.
-     * @param traitName Trait name to add.
-     * @param traitValue Trait value as a Node object.
-     */
-    public void onTrait(ShapeId target, String traitName, Node traitValue) {
-        onShapeTarget(traitName, traitValue.getSourceLocation(), id -> onTrait(target, id, traitValue));
-    }
-
     /**
      * Adds a trait to a shape.
      *
@@ -455,15 +297,23 @@ public void onTrait(ShapeId target, Trait trait) {
         pendingTraits.computeIfAbsent(target, targetId -> new ArrayList<>()).add(pending);
     }
 
+    /**
+     * Adds a forward reference that is resolved once all shapes have been loaded.
+     *
+     * @param expectedId The shape ID that would be resolved in the current namespace.
+     * @param consumer The consumer that receives the resolved shape ID.
+     */
+    void addForwardReference(ShapeId expectedId, Consumer<ShapeId> consumer) {
+        forwardReferenceResolvers.add(new ForwardReferenceResolver(expectedId, consumer));
+    }
+
     /**
      * Adds metadata to the loader.
      *
      * @param key Metadata key to add.
      * @param value Metadata value to add.
      */
     public void onMetadata(String key, Node value) {
-        validateState(value);
-
         if (!metadata.containsKey(key)) {
             metadata.put(key, value);
         } else if (metadata.get(key).isArrayNode() && value.isArrayNode()) {
@@ -486,23 +336,6 @@ public void onMetadata(String key, Node value) {
         }
     }
 
-    /**
-     * Registers an absolute shape ID as an alias.
-     *
-     * <p>These aliases are cleared when {@link #onOpenFile} is called.
-     *
-     * @param id Fully-qualified shape ID to register.
-     * @param location The location of where it is registered.
-     */
-    void onUseShape(ShapeId id, FromSourceLocation location) {
-        validateState(location);
-        ShapeId previous = useShapes.put(id.getName(), id);
-        if (previous != null) {
-            throw new UseException(String.format(
-                    "Cannot use name `%s` because it conflicts with `%s`", id, previous), location);
-        }
-    }
-
     /**
      * Called when the visitor has completed.
      *

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

@@ -84,7 +84,7 @@ public final class ModelAssembler {
     private ValidatorFactory validatorFactory;
 
     /**
-     * A map of files ot parse and load into the Model.
+     * A map of files to parse and load into the Model.
      *
      * <p>A {@code TreeMap} is used to ensure that JSON models are loaded
      * before IDL models. This is mostly a performance optimization. JSON

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

@@ -19,11 +19,13 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.not;
 
+import java.util.Optional;
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.SourceLocation;
 import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.ResourceShape;
+import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
 
 public class IdlModelLoaderTest {
@@ -77,4 +79,19 @@ public void canLoadAndAliasShapesAndTraits() {
                 .assemble()
                 .unwrap();
     }
+
+    @Test
+    public void defersApplyTargetAndTrait() {
+        Model model = Model.assembler()
+                .addImport(getClass().getResource("apply-use-1.smithy"))
+                .addImport(getClass().getResource("apply-use-2.smithy"))
+                .addImport(getClass().getResource("apply-use-3.smithy"))
+                .assemble()
+                .unwrap();
+
+        Shape shape = model.expectShape(ShapeId.from("smithy.example#Foo"));
+
+        assertThat(shape.findTrait(ShapeId.from("smithy.example#bar")), not(Optional.empty()));
+        assertThat(shape.findTrait(ShapeId.from("smithy.example.b#baz")), not(Optional.empty()));
+    }
 }

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

@@ -30,7 +30,6 @@
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.model.Model;
-import software.amazon.smithy.model.SourceLocation;
 import software.amazon.smithy.model.node.Node;
 import software.amazon.smithy.model.selector.Selector;
 import software.amazon.smithy.model.shapes.MemberShape;
@@ -39,6 +38,7 @@
 import software.amazon.smithy.model.shapes.StringShape;
 import software.amazon.smithy.model.shapes.StructureShape;
 import software.amazon.smithy.model.traits.DeprecatedTrait;
+import software.amazon.smithy.model.traits.DocumentationTrait;
 import software.amazon.smithy.model.traits.DynamicTrait;
 import software.amazon.smithy.model.traits.ReferencesTrait;
 import software.amazon.smithy.model.traits.TagsTrait;
@@ -52,15 +52,6 @@
 public class LoaderVisitorTest {
     private static final TraitFactory FACTORY = TraitFactory.createServiceFactory();
 
-    @Test
-    public void cannotMutateAfterOnEnd() {
-        Assertions.assertThrows(IllegalStateException.class, () -> {
-            LoaderVisitor visitor = new LoaderVisitor(FACTORY);
-            visitor.onEnd();
-            visitor.onMetadata("foo", Node.from("bar"));
-        });
-    }
-
     @Test
     public void callingOnEndTwiceIsIdempotent() {
         LoaderVisitor visitor = new LoaderVisitor(FACTORY);
@@ -128,12 +119,10 @@ public void cannotDuplicateNonPreludeTraitDefs() {
     @Test
     public void cannotDuplicateTraits() {
         LoaderVisitor visitor = new LoaderVisitor(FACTORY);
-        visitor.onOpenFile("/foo/baz");
-        visitor.onNamespace("foo.bam", SourceLocation.NONE);
         ShapeId id = ShapeId.from("foo.bam#Boo");
         visitor.onShape(StringShape.builder().id(id));
-        visitor.onTrait(id, "documentation", Node.from("abc"));
-        visitor.onTrait(id, "documentation", Node.from("def"));
+        visitor.onTrait(id, DocumentationTrait.ID, Node.from("abc"));
+        visitor.onTrait(id, DocumentationTrait.ID, Node.from("def"));
         List<ValidationEvent> events = visitor.onEnd().getValidationEvents();
 
         assertThat(events, not(empty()));
@@ -142,16 +131,14 @@ public void cannotDuplicateTraits() {
     @Test
     public void createsDynamicTraitWhenTraitFactoryReturnsEmpty() {
         LoaderVisitor visitor = new LoaderVisitor(FACTORY);
-        visitor.onOpenFile("/foo/baz.smithy");
-        visitor.onNamespace("foo.bam", SourceLocation.none());
         Shape def = StructureShape.builder()
                 .id("foo.baz#Bar")
                 .addTrait(TraitDefinition.builder().build())
                 .build();
         visitor.onShape(def);
         ShapeId id = ShapeId.from("foo.bam#Boo");
         visitor.onShape(StringShape.builder().id(id));
-        visitor.onTrait(id, "foo.baz#Bar", Node.from(true));
+        visitor.onTrait(id, ShapeId.from("foo.baz#Bar"), Node.from(true));
         Model model = visitor.onEnd().unwrap();
 
         assertThat(model.expectShape(id).findTrait("foo.baz#Bar").get(),
@@ -161,11 +148,9 @@ public void createsDynamicTraitWhenTraitFactoryReturnsEmpty() {
     @Test
     public void failsWhenTraitNotFound() {
         LoaderVisitor visitor = new LoaderVisitor(FACTORY);
-        visitor.onOpenFile("/foo/baz.smithy");
-        visitor.onNamespace("foo.bam", SourceLocation.none());
         ShapeId id = ShapeId.from("foo.bam#Boo");
         visitor.onShape(StringShape.builder().id(id));
-        visitor.onTrait(id, "foo.baz#Bar", Node.from(true));
+        visitor.onTrait(id, ShapeId.from("foo.baz#Bar"), Node.from(true));
         List<ValidationEvent> events = visitor.onEnd().getValidationEvents();
 
         assertThat(events, not(empty()));

smithy-model/src/test/resources/software/amazon/smithy/model/loader/apply-use-1.smithy

@@ -0,0 +1,9 @@
+namespace smithy.example
+
+use smithy.example.b#baz
+
+// This applies smithy.example#bar to smithy.example#Foo
+apply Foo @bar("hi")
+
+// This applies smithy.example.b#baz to smithy.example#Foo
+apply Foo @baz("hi")

smithy-model/src/test/resources/software/amazon/smithy/model/loader/apply-use-2.smithy

@@ -0,0 +1,6 @@
+namespace smithy.example
+
+string Foo
+
+@trait
+string bar

smithy-model/src/test/resources/software/amazon/smithy/model/loader/apply-use-3.smithy

@@ -0,0 +1,4 @@
+namespace smithy.example.b
+
+@trait
+string baz