Clean up model loading and add composite file

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

@@ -0,0 +1,128 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.model.loader;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.logging.Logger;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.shapes.ShapeType;
+import software.amazon.smithy.model.traits.Trait;
+import software.amazon.smithy.model.traits.TraitFactory;
+import software.amazon.smithy.model.validation.ValidationEvent;
+
+/**
+ * Aggregates together multiple {@code ModelFile}s.
+ */
+final class CompositeModelFile implements ModelFile {
+
+    private static final Logger LOGGER = Logger.getLogger(CompositeModelFile.class.getName());
+
+    private final TraitFactory traitFactory;
+    private final List<ModelFile> modelFiles;
+    private final List<ValidationEvent> mergeEvents = new ArrayList<>();
+
+    CompositeModelFile(TraitFactory traitFactory, List<ModelFile> modelFiles) {
+        this.traitFactory = traitFactory;
+        this.modelFiles = modelFiles;
+    }
+
+    @Override
+    public Set<ShapeId> shapeIds() {
+        Set<ShapeId> ids = new HashSet<>();
+        for (ModelFile modelFile : modelFiles) {
+            ids.addAll(modelFile.shapeIds());
+        }
+        return ids;
+    }
+
+    @Override
+    public ShapeType getShapeType(ShapeId id) {
+        for (ModelFile modFile : modelFiles) {
+            ShapeType fileType = modFile.getShapeType(id);
+            if (fileType != null) {
+                return fileType;
+            }
+        }
+        return null;
+    }
+
+    @Override
+    public Map<String, Node> metadata() {
+        MetadataContainer metadata = new MetadataContainer(mergeEvents);
+        for (ModelFile modelFile : modelFiles) {
+            for (Map.Entry<String, Node> entry : modelFile.metadata().entrySet()) {
+                metadata.putMetadata(entry.getKey(), entry.getValue());
+            }
+        }
+        return metadata.getData();
+    }
+
+    @Override
+    public TraitContainer resolveShapes(Set<ShapeId> ids, Function<ShapeId, ShapeType> typeProvider) {
+        TraitContainer.TraitHashMap traitValues = new TraitContainer.TraitHashMap(traitFactory, mergeEvents);
+        for (ModelFile modelFile : modelFiles) {
+            TraitContainer other = modelFile.resolveShapes(ids, typeProvider);
+            for (Map.Entry<ShapeId, Map<ShapeId, Trait>> entry : other.traits().entrySet()) {
+                ShapeId target = entry.getKey();
+                for (Map.Entry<ShapeId, Trait> appliedEntry : entry.getValue().entrySet()) {
+                    traitValues.onTrait(target, appliedEntry.getValue());
+                }
+            }
+        }
+        return traitValues;
+    }
+
+    @Override
+    public Collection<Shape> createShapes(TraitContainer resolvedTraits) {
+        // Merge all shapes together, resolve conflicts, and warn for acceptable conflicts.
+        Map<ShapeId, Shape> shapes = new HashMap<>();
+        for (ModelFile modelFile : modelFiles) {
+            for (Shape shape : modelFile.createShapes(resolvedTraits)) {
+                Shape previous = shapes.get(shape.getId());
+                if (previous == null) {
+                    shapes.put(shape.getId(), shape);
+                } else if (!previous.equals(shape)) {
+                    mergeEvents.add(LoaderUtils.onShapeConflict(shape.getId(), shape.getSourceLocation(),
+                                                                previous.getSourceLocation()));
+                } else {
+                    LOGGER.warning(() -> "Ignoring duplicate but equivalent shape definition: " + previous.getId()
+                                         + " defined at " + shape.getSourceLocation() + " and "
+                                         + previous.getSourceLocation());
+                }
+            }
+        }
+
+        return shapes.values();
+    }
+
+    @Override
+    public List<ValidationEvent> events() {
+        List<ValidationEvent> events = new ArrayList<>(mergeEvents);
+        for (ModelFile modelFile : modelFiles) {
+            events.addAll(modelFile.events());
+        }
+        return events;
+    }
+}

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

@@ -128,12 +128,12 @@ static final class TraitEntry {
         this.modelFile = new ForwardReferenceModelFile(traitFactory);
     }
 
-    List<ModelFile> parse() {
+    ModelFile parse() {
         ws();
         parseControlSection();
         parseMetadataSection();
         parseShapeSection();
-        return Collections.singletonList(modelFile);
+        return modelFile;
     }
 
     /**
@@ -696,8 +696,8 @@ private void onDeferredTrait(ShapeId target, String traitName, Node traitValue,
         });
     }
 
-    private Node coerceTraitValue(ShapeId traitId, Node value, boolean isAnnotation,
-            Function<ShapeId, ShapeType> typeProvider) {
+    private Node coerceTraitValue(
+            ShapeId traitId, Node value, boolean isAnnotation, Function<ShapeId, ShapeType> typeProvider) {
         if (isAnnotation && value.isNullNode()) {
             ShapeType targetType = typeProvider.apply(traitId);
             if (targetType != null) {

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

@@ -17,12 +17,10 @@
 
 import java.util.Collection;
 import java.util.List;
-import java.util.function.Function;
 import software.amazon.smithy.model.SourceLocation;
 import software.amazon.smithy.model.node.ExpectationNotMetException;
 import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ShapeId;
-import software.amazon.smithy.model.shapes.ShapeType;
 import software.amazon.smithy.model.validation.Severity;
 import software.amazon.smithy.model.validation.ValidationEvent;
 import software.amazon.smithy.model.validation.Validator;
@@ -82,25 +80,4 @@ static ValidationEvent onShapeConflict(ShapeId id, SourceLocation a, SourceLocat
                 .message(String.format("Conflicting shape definition for `%s` found at `%s` and `%s`", id, a, b))
                 .build();
     }
-
-    /**
-     * Iterates over ModelFiles to find the {@link ShapeType} of a shape.
-     *
-     * <p>The first found shape in a ModelFile wins. This is OK, since any
-     * kind of conflict is detected when the ModelFiles are merged together.
-     *
-     * @param modelFiles ModelFile instances to iterate over, searching for shapes by ID.
-     * @return Returns the found {@link ShapeType} or {@code null} if the shape does not exist.
-     */
-    public static Function<ShapeId, ShapeType> aggregateTypeProvider(List<ModelFile> modelFiles) {
-        return id -> {
-            for (ModelFile modFile : modelFiles) {
-                ShapeType fileType = modFile.getShapeType(id);
-                if (fileType != null) {
-                    return fileType;
-                }
-            }
-            return null;
-        };
-    }
 }

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

@@ -78,17 +78,6 @@ void putMetadata(String key, Node value) {
         }
     }
 
-    /**
-     * Merges another metadata container into this container.
-     *
-     * @param other Metadata container to merge into this container.
-     */
-    void mergeWith(Map<String, Node> other) {
-        for (Map.Entry<String, Node> entry : other.entrySet()) {
-            putMetadata(entry.getKey(), entry.getValue());
-        }
-    }
-
     /**
      * Gets all of the metadata in the container.
      *

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

@@ -28,12 +28,10 @@
 import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
-import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
@@ -513,7 +511,13 @@ public ValidatedResult<Model> assemble() {
         List<ModelFile> modelFiles = createModelFiles();
 
         try {
-            return doAssemble(modelFiles);
+            CompositeModelFile files = new CompositeModelFile(traitFactory, modelFiles);
+            TraitContainer traits = files.resolveShapes(files.shapeIds(), files::getShapeType);
+            Model model = Model.builder()
+                    .metadata(files.metadata())
+                    .addShapes(files.createShapes(traits))
+                    .build();
+            return validate(model, traits, files.events());
         } catch (SourceException e) {
             List<ValidationEvent> events = new ArrayList<>();
             events.add(ValidationEvent.fromSourceException(e));
@@ -554,21 +558,20 @@ private List<ModelFile> createModelFiles() {
         // Load parsed AST nodes and merge them into the assembler.
         for (Node node : documentNodes) {
             try {
-                modelFiles.addAll(ModelLoader.loadParsedNode(traitFactory, node));
+                modelFiles.add(ModelLoader.loadParsedNode(traitFactory, node));
             } catch (SourceException e) {
                 assemblerModelFile.events().add(ValidationEvent.fromSourceException(e));
             }
         }
 
         // Load model files and merge them into the assembler.
-        for (Map.Entry<String, Supplier<InputStream>> modelEntry : inputStreamModels.entrySet()) {
+        for (Map.Entry<String, Supplier<InputStream>> entry : inputStreamModels.entrySet()) {
             try {
-                List<ModelFile> loaded = ModelLoader.load(
-                        traitFactory, properties, modelEntry.getKey(), modelEntry.getValue());
-                if (loaded.isEmpty()) {
-                    LOGGER.warning(() -> "No ModelLoader was able to load " + modelEntry.getKey());
+                ModelFile loaded = ModelLoader.load(traitFactory, properties, entry.getKey(), entry.getValue());
+                if (loaded == null) {
+                    LOGGER.warning(() -> "No ModelLoader was able to load " + entry.getKey());
                 } else {
-                    modelFiles.addAll(loaded);
+                    modelFiles.add(loaded);
                 }
             } catch (SourceException e) {
                 assemblerModelFile.events().add(ValidationEvent.fromSourceException(e));
@@ -578,65 +581,29 @@ private List<ModelFile> createModelFiles() {
         return modelFiles;
     }
 
-    private ValidatedResult<Model> doAssemble(List<ModelFile> modelFiles) {
-        List<ValidationEvent> events = new ArrayList<>();
-        Set<ShapeId> ids = new HashSet<>();
-        MetadataContainer metadata = new MetadataContainer(events);
+    private ValidatedResult<Model> validate(Model model, TraitContainer traits, List<ValidationEvent> events) {
+        validateTraits(model.getShapeIds(), traits, events);
 
-        for (ModelFile modelFile : modelFiles) {
-            // Merge all metadata.
-            metadata.mergeWith(modelFile.metadata());
-            // Collect all known shape IDs across all modelFiles.
-            ids.addAll(modelFile.shapeIds());
-        }
-
-        Function<ShapeId, ShapeType> typeProvider = LoaderUtils.aggregateTypeProvider(modelFiles);
-
-        // Merge all pending traits across every model.
-        TraitContainer.TraitHashMap traitValues = new TraitContainer.TraitHashMap(traitFactory, events);
-        for (ModelFile modelFile : modelFiles) {
-            traitValues.merge(modelFile.resolveShapes(ids, typeProvider));
+        if (disableValidation) {
+            return new ValidatedResult<>(model, events);
         }
 
-        // Merge all shapes together, resolve conflicts, and warn for acceptable conflicts.
-        Map<ShapeId, Shape> shapes = new HashMap<>();
-        for (ModelFile modelFile : modelFiles) {
-            for (Shape shape : modelFile.createShapes(traitValues)) {
-                Shape previous = shapes.get(shape.getId());
-                if (previous == null) {
-                    shapes.put(shape.getId(), shape);
-                } else if (!previous.equals(shape)) {
-                    events.add(LoaderUtils.onShapeConflict(shape.getId(), shape.getSourceLocation(),
-                                                           previous.getSourceLocation()));
-                } else {
-                    LOGGER.warning(() -> "Ignoring duplicate but equivalent shape definition: " + previous.getId()
-                                         + " defined at " + shape.getSourceLocation() + " and "
-                                         + previous.getSourceLocation());
-                }
-            }
+        if (validatorFactory == null) {
+            validatorFactory = LazyValidatorFactoryHolder.INSTANCE;
         }
 
-        // Find traits applied to shapes that don't exist.
-        for (Map.Entry<ShapeId, Map<ShapeId, Trait>> entry : traitValues.traits().entrySet()) {
-            ShapeId target = entry.getKey();
-            if (!ids.contains(target)) {
-                for (Trait trait : entry.getValue().values()) {
-                    events.add(ValidationEvent.builder()
-                            .id(Validator.MODEL_ERROR)
-                            .severity(Severity.ERROR)
-                            .sourceLocation(trait)
-                            .message(String.format("Trait `%s` applied to unknown shape `%s`",
-                                                  Trait.getIdiomaticTraitName(trait.toShapeId()), target))
-                            .build());
-                }
-            }
-        }
+        // Validate the model based on the explicit validators and model metadata.
+        List<ValidationEvent> mergedEvents = ModelValidator.validate(model, validatorFactory, assembleValidators());
+        mergedEvents.addAll(events);
+        return new ValidatedResult<>(model, mergedEvents);
+    }
 
-        // Find trait values that weren't defined.
+    private void validateTraits(Set<ShapeId> ids, TraitContainer resolvedTraits, List<ValidationEvent> events) {
         Severity severity = areUnknownTraitsAllowed() ? Severity.WARNING : Severity.ERROR;
-        for (Map.Entry<ShapeId, Map<ShapeId, Trait>> entry : traitValues.traits().entrySet()) {
+        for (Map.Entry<ShapeId, Map<ShapeId, Trait>> entry : resolvedTraits.traits().entrySet()) {
             ShapeId target = entry.getKey();
             for (Trait trait : entry.getValue().values()) {
+                // Find trait values that weren't defined.
                 if (!ids.contains(trait.toShapeId())) {
                     events.add(ValidationEvent.builder()
                             .id(Validator.MODEL_ERROR)
@@ -648,41 +615,25 @@ private ValidatedResult<Model> doAssemble(List<ModelFile> modelFiles) {
                                     + "defined before it can be used in a model.", trait.toShapeId()))
                             .build());
                 }
+                // Find traits applied to shapes that don't exist.
+                if (!ids.contains(target)) {
+                    events.add(ValidationEvent.builder()
+                            .id(Validator.MODEL_ERROR)
+                            .severity(Severity.ERROR)
+                            .sourceLocation(trait)
+                            .message(String.format("Trait `%s` applied to unknown shape `%s`",
+                                                   Trait.getIdiomaticTraitName(trait.toShapeId()), target))
+                            .build());
+                }
             }
         }
-
-        for (ModelFile modelFile : modelFiles) {
-            events.addAll(modelFile.events());
-        }
-
-        Model.Builder builder = Model.builder();
-        builder.metadata(metadata.getData());
-        builder.addShapes(shapes.values());
-        Model model = builder.build();
-        return validate(model, events);
     }
 
     private boolean areUnknownTraitsAllowed() {
-        // Find trait values that weren't defined.
         Object allowUnknown = properties.get(ModelAssembler.ALLOW_UNKNOWN_TRAITS);
         return allowUnknown != null && (boolean) allowUnknown;
     }
 
-    private ValidatedResult<Model> validate(Model model, List<ValidationEvent> modelResultEvents) {
-        if (disableValidation) {
-            return new ValidatedResult<>(model, modelResultEvents);
-        }
-
-        if (validatorFactory == null) {
-            validatorFactory = LazyValidatorFactoryHolder.INSTANCE;
-        }
-
-        // Validate the model based on the explicit validators and model metadata.
-        List<ValidationEvent> events = ModelValidator.validate(model, validatorFactory, assembleValidators());
-        events.addAll(modelResultEvents);
-        return new ValidatedResult<>(model, events);
-    }
-
     private List<Validator> assembleValidators() {
         // Find and register built-in validators with the validator.
         List<Validator> copiedValidators = new ArrayList<>(validatorFactory.loadBuiltinValidators());