Make various perf improvements

This commit makes various performance improvements that particularly help
when validating models with 100k+ shapes.

* Fixed JMH benchmarks so they work with the updated Gradle version.
* Model now uses a synchronized IdentityHashMap now to back the
  blackboard cache rather than a ConcurrentHashMap. This will actually
  now prevent duplicate work in creating KnowledgeIndexes. An
  IdentityHashMap was used because it works well for the classes
  used to cache knowledge indexes.
* HttpBindingIndex now uses a WeakReference to Model. This was
  previously an unnecessary cyclical reference.
* HttpBindingIndex no longer throws when attempting to access the
  HTTP bindings of shapes that don't exist or aren't operations.
  This prevents it from having to store a List entry for every
  single operation.
* Model#getShapesWithTraits was used everywhere possible rather
  than streaming over all shapes and looking for traits.
* Made optimizations to NullableIndex to no longer need to
  traverse every shape.
* Removed unnecessary streaming from OperationIndex.
* Removed unnecessary streaming from PaginatedIndex.
* Removed unnecessary streaming from NeighborVisitor.
* Updated Node expect methods to also accept a Supplier to create
  error messages if the expectation fails. This prevents needing to
  evaluate String.format even for valid node assertions.
* AttributeSelector no longer uses a BiFunction key supplier, and instead
  the attribute path is just passed in. This allows for the selector
  to also perform optimizations when determining if a shape has a trait
  by leveraging Model#getShapesWithTraits.
* InternalSelectors used to implement Selectors now support more general
  optimizations. This was previously hardcoded to only support an optimization
  for selecting shapes by type, but now selecting shapes by trait is
  optimized too.
* Minor optimization to structure and union loading so that when validating
  that members have a shape ID compatible with the container, an intermediate
  shape ID no longer is constructed.
* The ShapeId cache was increased from 1024 to 8192. This helps significantly
  with large models. The ShapeId cache was also updated to implement the
  LRA cache inside of the computeIfAbsent method.
* NodeValidationVisitor now has a Context object that supports caching
  selectors evaluated against a model. This helps significantly with IdRef
  validation.  To make this caching reusable, the visitor is now mutable
  after it is constructed.
* NodeValidationVisitor idRef now special cases "*" and uses the context
  cache.
* TraitTargetValidator has been simplified, special cases "*", and now
  uses a cache to speed up evaluating traits that use the same selectors.
* TraitValueValidator now reuses the same NodeValidationVisitor in order
  to reuse the same selector cache.
* The `Model#toSet(Class)` method is now public, so getting a set of
  shapes of a specific type no longer has to always go through a
  `Stream`.

build.gradle

@@ -23,7 +23,7 @@ plugins {
     id "jacoco"
     id "com.github.spotbugs" version "4.6.0"
     id "io.codearte.nexus-staging" version "0.21.0"
-    id "me.champeau.gradle.jmh" version "0.5.3"
+    id "me.champeau.jmh" version "0.6.4"
 }
 
 ext {

config/spotbugs/filter.xml

@@ -104,4 +104,11 @@
         <Class name="software.amazon.smithy.model.knowledge.NeighborProviderIndex"/>
         <Bug pattern="NP_NULL_ON_SOME_PATH_FROM_RETURN_VALUE"/>
     </Match>
+
+    <!-- Spotbugs for some reason isn't seeing that Objects.requireNonNull prevents a null return.
+         this is used when dereferencing a WeakReference to the Model. -->
+    <Match>
+        <Class name="software.amazon.smithy.model.knowledge.HttpBindingIndex"/>
+        <Bug pattern="NP_NULL_ON_SOME_PATH_FROM_RETURN_VALUE"/>
+    </Match>
 </FindBugsFilter>

smithy-model/build.gradle

@@ -21,7 +21,7 @@ ext {
     moduleName = "software.amazon.smithy.model"
 }
 
-apply plugin: "me.champeau.gradle.jmh"
+apply plugin: "me.champeau.jmh"
 
 dependencies {
     api project(":smithy-utils")

smithy-model/src/jmh/resources/software/amazon/smithy/model/jmh/test-model.smithy

@@ -18,7 +18,6 @@ operation Resource1Operation {}
 resource Resource1_2 {}
 
 resource Resource1_1 {
-    type: resource,
     operations: [Resource1_1_Operation],
     resources: [Resource1_1_1, Resource1_1_2]
 }

smithy-model/src/main/java/software/amazon/smithy/model/Model.java

@@ -20,6 +20,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.IdentityHashMap;
 import java.util.Iterator;
 import java.util.Map;
 import java.util.Objects;
@@ -64,7 +65,8 @@ public final class Model implements ToSmithyBuilder<Model> {
     private final Map<Class<? extends Shape>, Set<? extends Shape>> cachedTypes = new ConcurrentHashMap<>();
 
     /** Cache of computed {@link KnowledgeIndex} instances. */
-    private final Map<Class<? extends KnowledgeIndex>, KnowledgeIndex> blackboard = new ConcurrentHashMap<>();
+    private final Map<Class<? extends KnowledgeIndex>, KnowledgeIndex> blackboard
+            = Collections.synchronizedMap(new IdentityHashMap<>());
 
     /** Lazily computed trait mappings. */
     private volatile TraitCache traitCache;
@@ -283,7 +285,7 @@ public <T extends Shape> Stream<T> shapes(Class<T> shapeType) {
      * @return Returns an unmodifiable set of shapes.
      */
     @SuppressWarnings("unchecked")
-    private <T extends Shape> Set<T> toSet(Class<T> shapeType) {
+    public <T extends Shape> Set<T> toSet(Class<T> shapeType) {
         return (Set<T>) cachedTypes.computeIfAbsent(shapeType, t -> {
             Set<T> result = new HashSet<>();
             for (Shape shape : shapeMap.values()) {
@@ -394,22 +396,7 @@ public <T extends KnowledgeIndex> T getKnowledge(Class<T> type) {
      */
     @SuppressWarnings("unchecked")
     public <T extends KnowledgeIndex> T getKnowledge(Class<T> type, Function<Model, T> constructor) {
-        // This method intentionally does not use putIfAbsent to avoid
-        // deadlocks in the case where a knowledge index needs to access
-        // other knowledge indexes from a Model. While this *can* cause
-        // duplicate computations, it's preferable over *always* requiring
-        // duplicate computations, deadlocks of computeIfAbsent, or
-        // spreading out the cache state associated with previously
-        // computed indexes through WeakHashMaps on each KnowledgeIndex
-        // (a previous approach we used for caching).
-        T value = (T) blackboard.get(type);
-
-        if (value == null) {
-            value = constructor.apply(this);
-            blackboard.put(type, value);
-        }
-
-        return value;
+        return (T) blackboard.computeIfAbsent(type, t -> constructor.apply(this));
     }
 
     /**

smithy-model/src/main/java/software/amazon/smithy/model/knowledge/BottomUpIndex.java

@@ -40,7 +40,8 @@ public final class BottomUpIndex implements KnowledgeIndex {
 
     public BottomUpIndex(Model model) {
         PathFinder pathFinder = PathFinder.create(model);
-        model.shapes(ServiceShape.class).forEach(service -> {
+
+        for (ServiceShape service : model.toSet(ServiceShape.class)) {
             Map<ShapeId, List<EntityShape>> serviceBindings = new HashMap<>();
             parentBindings.put(service.getId(), serviceBindings);
             for (PathFinder.Path path : pathFinder.search(service, SELECTOR)) {
@@ -56,7 +57,7 @@ public BottomUpIndex(Model model) {
                 // Add the end shape (a resource or operation) to the service bindings.
                 serviceBindings.put(path.getEndShape().getId(), shapes);
             }
-        });
+        }
     }
 
     public static BottomUpIndex of(Model model) {

smithy-model/src/main/java/software/amazon/smithy/model/knowledge/EventStreamIndex.java

@@ -45,14 +45,14 @@ public final class EventStreamIndex implements KnowledgeIndex {
     public EventStreamIndex(Model model) {
         OperationIndex operationIndex = OperationIndex.of(model);
 
-        model.shapes(OperationShape.class).forEach(operation -> {
+        for (OperationShape operation : model.toSet(OperationShape.class)) {
             operationIndex.getInput(operation).ifPresent(input -> {
                 computeEvents(model, operation, input, inputInfo);
             });
             operationIndex.getOutput(operation).ifPresent(output -> {
                 computeEvents(model, operation, output, outputInfo);
             });
-        });
+        }
     }
 
     public static EventStreamIndex of(Model model) {

smithy-model/src/main/java/software/amazon/smithy/model/knowledge/HttpBindingIndex.java

@@ -15,12 +15,14 @@
 
 package software.amazon.smithy.model.knowledge;
 
+import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -44,8 +46,6 @@
 import software.amazon.smithy.model.traits.MediaTypeTrait;
 import software.amazon.smithy.model.traits.StreamingTrait;
 import software.amazon.smithy.model.traits.TimestampFormatTrait;
-import software.amazon.smithy.model.traits.Trait;
-import software.amazon.smithy.utils.ListUtils;
 
 /**
  * Computes and indexes the explicit and implicit HTTP bindings of a model.
@@ -59,29 +59,27 @@
  * <p>This index does not perform validation of the underlying model.
  */
 public final class HttpBindingIndex implements KnowledgeIndex {
-    private final Model model;
+    private final WeakReference<Model> model;
     private final Map<ShapeId, List<HttpBinding>> requestBindings = new HashMap<>();
     private final Map<ShapeId, List<HttpBinding>> responseBindings = new HashMap<>();
 
     public HttpBindingIndex(Model model) {
-        this.model = model;
+        this.model = new WeakReference<>(model);
         OperationIndex opIndex = OperationIndex.of(model);
-        model.shapes(OperationShape.class).forEach(shape -> {
-            if (shape.getTrait(HttpTrait.class).isPresent()) {
-                requestBindings.put(shape.getId(), computeRequestBindings(opIndex, shape));
-                responseBindings.put(shape.getId(), computeResponseBindings(opIndex, shape));
-            } else {
-                requestBindings.put(shape.getId(), ListUtils.of());
-                responseBindings.put(shape.getId(), ListUtils.of());
-            }
-        });
+
+        for (Shape shape : model.getShapesWithTrait(HttpTrait.class)) {
+            shape.asOperationShape().ifPresent(operation -> {
+                requestBindings.put(operation.getId(), computeRequestBindings(opIndex, operation));
+                responseBindings.put(operation.getId(), computeResponseBindings(opIndex, operation));
+            });
+        }
 
         // Add error structure bindings.
-        model.shapes(StructureShape.class)
-                .flatMap(shape -> Trait.flatMapStream(shape, ErrorTrait.class))
-                .forEach(pair -> responseBindings.put(
-                        pair.getLeft().getId(),
-                        createStructureBindings(pair.getLeft(), false)));
+        for (Shape shape : model.getShapesWithTrait(ErrorTrait.class)) {
+            shape.asStructureShape().ifPresent(structure -> {
+                responseBindings.put(structure.getId(), createStructureBindings(structure, false));
+            });
+        }
     }
 
     public static HttpBindingIndex of(Model model) {
@@ -120,14 +118,18 @@ public static boolean hasHttpResponseBindings(Shape shape) {
 
     private HttpTrait getHttpTrait(ToShapeId operation) {
         ShapeId id = operation.toShapeId();
-        return model.getShape(id)
+        return getModel().getShape(id)
                 .orElseThrow(() -> new IllegalArgumentException(id + " is not a valid shape"))
                 .asOperationShape()
                 .orElseThrow(() -> new IllegalArgumentException(id + " is not an operation shape"))
                 .getTrait(HttpTrait.class)
                 .orElseThrow(() -> new IllegalArgumentException(id + " has no http binding trait"));
     }
 
+    private Model getModel() {
+        return Objects.requireNonNull(model.get(), "The dereferenced WeakReference<Model> is null");
+    }
+
     /**
      * Gets the computed status code of an operation or error structure.
      *
@@ -138,7 +140,7 @@ private HttpTrait getHttpTrait(ToShapeId operation) {
      */
     public int getResponseCode(ToShapeId shapeOrId) {
         ShapeId id = shapeOrId.toShapeId();
-        Shape shape = model.getShape(id).orElseThrow(() -> new IllegalArgumentException("Shape not found " + id));
+        Shape shape = getModel().getShape(id).orElseThrow(() -> new IllegalArgumentException("Shape not found " + id));
 
         if (shape.isOperationShape()) {
             return getHttpTrait(id).getCode();
@@ -157,34 +159,26 @@ public int getResponseCode(ToShapeId shapeOrId) {
      *
      * @param operationShapeOrId Operation to get the request bindings for.
      * @return Map of unmodifiable bindings.
-     * @throws IllegalArgumentException if the given shape is not an operation.
      */
     public Map<String, HttpBinding> getRequestBindings(ToShapeId operationShapeOrId) {
         ShapeId id = operationShapeOrId.toShapeId();
-        validateRequestBindingShapeId(id);
-        return requestBindings.get(id).stream()
+        return requestBindings.getOrDefault(id, Collections.emptyList())
+                .stream()
                 .collect(Collectors.toMap(HttpBinding::getMemberName, Function.identity()));
     }
 
-    private void validateRequestBindingShapeId(ShapeId id) {
-        if (!requestBindings.containsKey(id)) {
-            throw new IllegalArgumentException(id + " does not reference an operation with http bindings");
-        }
-    }
-
     /**
      * Gets the request bindings of an operation as a map of member name to
      * the binding for a specific location type.
      *
      * @param operationShapeOrId Operation to get the request bindings for.
      * @param requestLocation Location of the binding.
      * @return Map of unmodifiable bindings.
-     * @throws IllegalArgumentException if the given shape is not an operation.
      */
     public List<HttpBinding> getRequestBindings(ToShapeId operationShapeOrId, HttpBinding.Location requestLocation) {
         ShapeId id = operationShapeOrId.toShapeId();
-        validateRequestBindingShapeId(id);
-        return requestBindings.get(id).stream()
+        return requestBindings.getOrDefault(id, Collections.emptyList())
+                .stream()
                 .filter(binding -> binding.getLocation() == requestLocation)
                 .collect(Collectors.toList());
     }
@@ -195,22 +189,14 @@ public List<HttpBinding> getRequestBindings(ToShapeId operationShapeOrId, HttpBi
      *
      * @param shapeOrId Operation or error structure shape or ID.
      * @return Map of unmodifiable bindings.
-     * @throws IllegalArgumentException if the given shape is not an operation
-     *  or error structure.
      */
     public Map<String, HttpBinding> getResponseBindings(ToShapeId shapeOrId) {
         ShapeId id = shapeOrId.toShapeId();
-        validateResponseBindingShapeId(id);
-        return responseBindings.get(id).stream()
+        return responseBindings.getOrDefault(id, Collections.emptyList())
+                .stream()
                 .collect(Collectors.toMap(HttpBinding::getMemberName, Function.identity()));
     }
 
-    private void validateResponseBindingShapeId(ShapeId id) {
-        if (!responseBindings.containsKey(id)) {
-            throw new IllegalArgumentException(id + " does not reference an operation or error structure");
-        }
-    }
-
     /**
      * Gets the computed HTTP message response bindings for an operation
      * or structure with an error trait for a specific binding type.
@@ -223,8 +209,8 @@ private void validateResponseBindingShapeId(ShapeId id) {
      */
     public List<HttpBinding> getResponseBindings(ToShapeId shapeOrId, HttpBinding.Location bindingLocation) {
         ShapeId id = shapeOrId.toShapeId();
-        validateResponseBindingShapeId(id);
-        return responseBindings.get(id).stream()
+        return responseBindings.getOrDefault(id, Collections.emptyList())
+                .stream()
                 .filter(binding -> binding.getLocation() == bindingLocation)
                 .collect(Collectors.toList());
     }
@@ -243,6 +229,7 @@ public TimestampFormatTrait.Format determineTimestampFormat(
             HttpBinding.Location location,
             TimestampFormatTrait.Format defaultFormat
     ) {
+        Model model = getModel();
         return model.getShape(member.toShapeId())
                 // Use the timestampFormat trait on the member or target if present.
                 .flatMap(shape -> shape.getMemberTrait(model, TimestampFormatTrait.class))
@@ -371,6 +358,8 @@ private String determineContentType(
             String documentContentType,
             String eventStreamContentType
     ) {
+        Model model = getModel();
+
         for (HttpBinding binding : bindings) {
             if (binding.getLocation() == HttpBinding.Location.DOCUMENT) {
                 return documentContentType;

smithy-model/src/main/java/software/amazon/smithy/model/knowledge/IdentifierBindingIndex.java

@@ -53,7 +53,9 @@ public enum BindingType {
 
     public IdentifierBindingIndex(Model model) {
         OperationIndex operationIndex = OperationIndex.of(model);
-        model.shapes(ResourceShape.class).forEach(resource -> processResource(resource, operationIndex, model));
+        for (ResourceShape resource : model.toSet(ResourceShape.class)) {
+            processResource(resource, operationIndex, model);
+        }
     }
 
     public static IdentifierBindingIndex of(Model model) {

smithy-model/src/main/java/software/amazon/smithy/model/knowledge/NullableIndex.java

@@ -48,15 +48,23 @@ public class NullableIndex implements KnowledgeIndex {
     private final Set<ShapeId> nullableShapes = new HashSet<>();
 
     public NullableIndex(Model model) {
-        for (Shape shape : model.toSet()) {
-            if (shape.asMemberShape().isPresent()) {
-                if (isMemberNullable(model, shape.asMemberShape().get())) {
-                    nullableShapes.add(shape.getId());
-                }
-            } else if (isShapeBoxed(shape)) {
+        for (ShapeType type : INHERENTLY_BOXED) {
+            model.shapes(type.getShapeClass()).forEach(shape -> nullableShapes.add(shape.getId()));
+        }
+
+        for (Shape shape : model.getShapesWithTrait(BoxTrait.class)) {
+            // Only structure members honor the box trait, so defer to the
+            // isMemberNullable method to determine member nullability.
+            if (!shape.isMemberShape()) {
                 nullableShapes.add(shape.getId());
             }
         }
+
+        for (MemberShape member : model.toSet(MemberShape.class)) {
+            if (isMemberNullable(model, member)) {
+                nullableShapes.add(member.getId());
+            }
+        }
     }
 
     public static NullableIndex of(Model model) {

smithy-model/src/main/java/software/amazon/smithy/model/knowledge/OperationIndex.java

@@ -15,12 +15,12 @@
 
 package software.amazon.smithy.model.knowledge;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.stream.Collectors;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.OperationShape;
@@ -43,21 +43,20 @@ public final class OperationIndex implements KnowledgeIndex {
     private final Map<ShapeId, List<StructureShape>> errors = new HashMap<>();
 
     public OperationIndex(Model model) {
-        model.shapes(OperationShape.class).forEach(operation -> {
+        for (OperationShape operation : model.toSet(OperationShape.class)) {
             operation.getInput()
                     .flatMap(id -> getStructure(model, id))
                     .ifPresent(shape -> inputs.put(operation.getId(), shape));
             operation.getOutput()
                     .flatMap(id -> getStructure(model, id))
                     .ifPresent(shape -> outputs.put(operation.getId(), shape));
-            errors.put(operation.getId(),
-                       operation.getErrors()
-                               .stream()
-                               .map(e -> getStructure(model, e))
-                               .filter(Optional::isPresent)
-                               .map(Optional::get)
-                               .collect(Collectors.toList()));
-        });
+
+            List<StructureShape> errorShapes = new ArrayList<>(operation.getErrors().size());
+            for (ShapeId target : operation.getErrors()) {
+                model.getShape(target).flatMap(Shape::asStructureShape).ifPresent(errorShapes::add);
+            }
+            errors.put(operation.getId(), errorShapes);
+        }
     }
 
     public static OperationIndex of(Model model) {