Make various perf improvements

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/checkstyle/checkstyle.xml

@@ -171,7 +171,6 @@
         <!-- See http://checkstyle.sf.net/config_coding.html -->
         <module name="EqualsHashCode"/>
         <module name="IllegalInstantiation"/>
-        <module name="InnerAssignment"/>
         <module name="MissingSwitchDefault"/>
         <module name="SimplifyBooleanExpression"/>
         <module name="SimplifyBooleanReturn"/>

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-aws-traits/src/main/java/software/amazon/smithy/aws/traits/ArnIndex.java

@@ -32,25 +32,26 @@
 import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.ResourceShape;
 import software.amazon.smithy.model.shapes.ServiceShape;
+import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.ToShapeId;
-import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.utils.Pair;
 
 /**
  * Resolves and indexes the ARN templates for each resource in a service.
  */
 public final class ArnIndex implements KnowledgeIndex {
-    private final Map<ShapeId, String> arnServices;
+    private final Map<ShapeId, String> arnServices = new HashMap<>();
     private final Map<ShapeId, Map<ShapeId, ArnTrait>> templates;
     private final Map<ShapeId, Map<ShapeId, ArnTrait>> effectiveArns = new HashMap<>();
 
     public ArnIndex(Model model) {
         // Pre-compute the ARN services.
-        arnServices = unmodifiableMap(model.shapes(ServiceShape.class)
-                .flatMap(shape -> Trait.flatMapStream(shape, ServiceTrait.class))
-                .map(pair -> Pair.of(pair.getLeft().getId(), resolveServiceArn(pair)))
-                .collect(Collectors.toMap(Pair::getLeft, Pair::getRight)));
+        for (Shape shape : model.getShapesWithTrait(ServiceTrait.class)) {
+            shape.asServiceShape().ifPresent(service -> {
+                arnServices.put(service.getId(), service.expectTrait(ServiceTrait.class).getArnNamespace());
+            });
+        }
 
         // Pre-compute all of the ArnTemplates in a service shape.
         TopDownIndex topDownIndex = TopDownIndex.of(model);
@@ -73,15 +74,15 @@ public static ArnIndex of(Model model) {
         return model.getKnowledge(ArnIndex.class, ArnIndex::new);
     }
 
-    private static String resolveServiceArn(Pair<ServiceShape, ServiceTrait> pair) {
-        return pair.getRight().getArnNamespace();
-    }
-
     private Pair<ShapeId, Map<ShapeId, ArnTrait>> compileServiceArns(TopDownIndex index, ServiceShape service) {
-        return Pair.of(service.getId(), unmodifiableMap(index.getContainedResources(service.getId()).stream()
-                .flatMap(resource -> Trait.flatMapStream(resource, ArnTrait.class))
-                .map(pair -> Pair.of(pair.getLeft().getId(), pair.getRight()))
-                .collect(Collectors.toMap(Pair::getLeft, Pair::getRight))));
+        Map<ShapeId, ArnTrait> mapping = new HashMap<>();
+        for (ResourceShape resource : index.getContainedResources(service.getId())) {
+            resource.getTrait(ArnTrait.class).ifPresent(arnTrait -> {
+                mapping.put(resource.getId(), arnTrait);
+            });
+        }
+
+        return Pair.of(service.getId(), Collections.unmodifiableMap(mapping));
     }
 
     private void compileEffectiveArns(

smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/ArnTemplateValidator.java

@@ -18,22 +18,22 @@
 import static java.util.stream.Collectors.toList;
 import static software.amazon.smithy.model.validation.ValidationUtils.tickedList;
 
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.regex.Pattern;
-import java.util.stream.Stream;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.shapes.ResourceShape;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.Shape;
-import software.amazon.smithy.model.traits.Trait;
+import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.validation.AbstractValidator;
 import software.amazon.smithy.model.validation.ValidationEvent;
-import software.amazon.smithy.utils.OptionalUtils;
-import software.amazon.smithy.utils.Pair;
 import software.amazon.smithy.utils.SmithyInternalApi;
 
 /**
@@ -47,28 +47,33 @@ public final class ArnTemplateValidator extends AbstractValidator {
     @Override
     public List<ValidationEvent> validate(Model model) {
         ArnIndex arnIndex = ArnIndex.of(model);
-        return model.shapes(ServiceShape.class)
-                .flatMap(service -> Trait.flatMapStream(service, ServiceTrait.class))
-                .flatMap(pair -> validateService(model, arnIndex, pair.getLeft()))
-                .collect(toList());
+        List<ValidationEvent> events = new ArrayList<>();
+        for (Shape shape : model.getShapesWithTrait(ServiceTrait.class)) {
+            shape.asServiceShape().ifPresent(service -> events.addAll(validateService(model, arnIndex, service)));
+        }
+        return events;
     }
 
-    private Stream<ValidationEvent> validateService(Model model, ArnIndex arnIndex, ServiceShape service) {
+    private List<ValidationEvent> validateService(Model model, ArnIndex arnIndex, ServiceShape service) {
         // Make sure each ARN template contains relevant identifiers.
-        return arnIndex.getServiceResourceArns(service.getId()).entrySet().stream()
-                .flatMap(entry -> OptionalUtils.stream(model.getShape(entry.getKey())
-                        .flatMap(Shape::asResourceShape)
-                        .map(resource -> Pair.of(resource, entry.getValue()))))
-                .flatMap(pair -> validateResourceArn(pair.getLeft(), pair.getRight()));
+        List<ValidationEvent> events = new ArrayList<>();
+        for (Map.Entry<ShapeId, ArnTrait> entry : arnIndex.getServiceResourceArns(service.getId()).entrySet()) {
+            model.getShape(entry.getKey()).flatMap(Shape::asResourceShape).ifPresent(resource -> {
+                events.addAll(validateResourceArn(resource, entry.getValue()));
+            });
+        }
+        return events;
     }
 
-    private Stream<ValidationEvent> validateResourceArn(ResourceShape resource, ArnTrait template) {
+    private List<ValidationEvent> validateResourceArn(ResourceShape resource, ArnTrait template) {
         // Fail early on syntax error, otherwise, validate that the
         // template correspond to identifiers.
-        return syntax(resource, template).map(Stream::of).orElseGet(() -> Stream.concat(
-                OptionalUtils.stream(enough(resource.getIdentifiers().keySet(), resource, template)),
-                OptionalUtils.stream(tooMuch(resource.getIdentifiers().keySet(), resource, template)))
-        );
+        return syntax(resource, template).map(Collections::singletonList).orElseGet(() -> {
+            List<ValidationEvent> events = new ArrayList<>();
+            enough(resource.getIdentifiers().keySet(), resource, template).ifPresent(events::add);
+            tooMuch(resource.getIdentifiers().keySet(), resource, template).ifPresent(events::add);
+            return events;
+        });
     }
 
     // Validates the syntax of each template.

smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/EventSourceValidator.java

@@ -15,19 +15,17 @@
 
 package software.amazon.smithy.aws.traits;
 
-import static java.util.stream.Collectors.toList;
-
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.shapes.ServiceShape;
-import software.amazon.smithy.model.traits.Trait;
+import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.validation.AbstractValidator;
 import software.amazon.smithy.model.validation.ValidationEvent;
 import software.amazon.smithy.utils.MapUtils;
-import software.amazon.smithy.utils.OptionalUtils;
 import software.amazon.smithy.utils.SmithyInternalApi;
 import software.amazon.smithy.utils.StringUtils;
 
@@ -45,10 +43,13 @@ public final class EventSourceValidator extends AbstractValidator {
 
     @Override
     public List<ValidationEvent> validate(Model model) {
-        return model.shapes(ServiceShape.class)
-                .flatMap(service -> Trait.flatMapStream(service, ServiceTrait.class))
-                .flatMap(pair -> OptionalUtils.stream(validateService(pair.getLeft(), pair.getRight())))
-                .collect(toList());
+        List<ValidationEvent> events = new ArrayList<>();
+        for (Shape shape : model.getShapesWithTrait(ServiceTrait.class)) {
+            shape.asServiceShape()
+                    .flatMap(service -> validateService(service, service.expectTrait(ServiceTrait.class)))
+                    .ifPresent(events::add);
+        }
+        return events;
     }
 
     private Optional<ValidationEvent> validateService(ServiceShape service, ServiceTrait trait) {

smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/SdkServiceIdValidator.java

@@ -15,8 +15,6 @@
 
 package software.amazon.smithy.aws.traits;
 
-import static java.util.stream.Collectors.toList;
-
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Locale;
@@ -25,11 +23,10 @@
 import java.util.regex.Pattern;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.shapes.ServiceShape;
-import software.amazon.smithy.model.traits.Trait;
+import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.validation.AbstractValidator;
 import software.amazon.smithy.model.validation.ValidationEvent;
 import software.amazon.smithy.model.validation.ValidationUtils;
-import software.amazon.smithy.utils.OptionalUtils;
 import software.amazon.smithy.utils.SetUtils;
 import software.amazon.smithy.utils.SmithyInternalApi;
 
@@ -73,10 +70,13 @@ public final class SdkServiceIdValidator extends AbstractValidator {
 
     @Override
     public List<ValidationEvent> validate(Model model) {
-        return model.shapes(ServiceShape.class)
-                .flatMap(service -> Trait.flatMapStream(service, ServiceTrait.class))
-                .flatMap(pair -> OptionalUtils.stream(validateService(pair.getLeft(), pair.getRight())))
-                .collect(toList());
+        List<ValidationEvent> events = new ArrayList<>();
+        for (Shape shape : model.getShapesWithTrait(ServiceTrait.class)) {
+            shape.asServiceShape()
+                    .flatMap(service -> validateService(service, service.expectTrait(ServiceTrait.class)))
+                    .ifPresent(events::add);
+        }
+        return events;
     }
 
     /**

smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/clientendpointdiscovery/CleanClientDiscoveryTraitTransformer.java

@@ -27,7 +27,6 @@
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.ErrorTrait;
-import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.model.transform.ModelTransformer;
 import software.amazon.smithy.model.transform.ModelTransformerPlugin;
 import software.amazon.smithy.utils.SmithyInternalApi;
@@ -62,16 +61,18 @@ public Model onRemove(ModelTransformer transformer, Collection<Shape> shapes, Mo
     }
 
     private Set<Shape> getServicesToUpdate(Model model, Set<ShapeId> removedOperations, Set<ShapeId> removedErrors) {
-        return model.shapes(ServiceShape.class)
-                .flatMap(service -> Trait.flatMapStream(service, ClientEndpointDiscoveryTrait.class))
-                .filter(pair -> removedOperations.contains(pair.getRight().getOperation())
-                        || removedErrors.contains(pair.getRight().getError()))
-                .map(pair -> {
-                    ServiceShape.Builder builder = pair.getLeft().toBuilder();
+        Set<Shape> result = new HashSet<>();
+        for (Shape shape : model.getShapesWithTrait(ClientEndpointDiscoveryTrait.class)) {
+            shape.asServiceShape().ifPresent(service -> {
+                ClientEndpointDiscoveryTrait trait = service.expectTrait(ClientEndpointDiscoveryTrait.class);
+                if (removedOperations.contains(trait.getOperation()) || removedErrors.contains(trait.getError())) {
+                    ServiceShape.Builder builder = service.toBuilder();
                     builder.removeTrait(ClientEndpointDiscoveryTrait.ID);
-                    return builder.build();
-                })
-                .collect(Collectors.toSet());
+                    result.add(builder.build());
+                }
+            });
+        }
+        return result;
     }
 
     private Set<Shape> getOperationsToUpdate(
@@ -89,16 +90,19 @@ private Set<Shape> getOperationsToUpdate(
                 .flatMap(service -> discoveryIndex.getEndpointDiscoveryOperations(service).stream())
                 .collect(Collectors.toSet());
 
-        return model.shapes(OperationShape.class)
-                // Get all endpoint discovery operations
-                .flatMap(operation -> Trait.flatMapStream(operation, ClientDiscoveredEndpointTrait.class))
-                // Only get the ones where discovery is optional, as it is safe to remove in that case
-                .filter(pair -> !pair.getRight().isRequired())
-                // Only get the ones that aren't still bound to a service that requires endpoint discovery
-                .filter(pair -> !stillBoundOperations.contains(pair.getLeft().getId()))
-                // Remove the trait
-                .map(pair -> pair.getLeft().toBuilder().removeTrait(ClientDiscoveredEndpointTrait.ID).build())
-                .collect(Collectors.toSet());
+        // Get all endpoint discovery operations
+        Set<Shape> result = new HashSet<>();
+        for (Shape shape : model.getShapesWithTrait(ClientDiscoveredEndpointTrait.class)) {
+            shape.asOperationShape().ifPresent(operation -> {
+                ClientDiscoveredEndpointTrait trait = operation.expectTrait(ClientDiscoveredEndpointTrait.class);
+                // Only get the ones where discovery is optional, as it is safe to remove in that case.
+                // Only get the ones that aren't still bound to a service that requires endpoint discovery.
+                if (!trait.isRequired() && !stillBoundOperations.contains(operation.getId())) {
+                    result.add(operation.toBuilder().removeTrait(ClientDiscoveredEndpointTrait.ID).build());
+                }
+            });
+        }
+        return result;
     }
 
     private Set<Shape> getMembersToUpdate(Model model, Set<ShapeId> updatedOperations) {