Improve TraitTargetValidator

smithy-model/src/main/java/software/amazon/smithy/model/selector/Selector.java

@@ -27,8 +27,9 @@
  */
 @FunctionalInterface
 public interface Selector {
+
     /** A selector that always returns all provided values. */
-    Selector IDENTITY = (model, visitor, shapes) -> shapes;
+    Selector IDENTITY = new WrappedSelector("*", (model, provider, shapes) -> shapes);
 
     /**
      * Matches a selector to a set of shapes.

smithy-model/src/main/java/software/amazon/smithy/model/validation/validators/TraitTargetValidator.java

@@ -15,14 +15,19 @@
 
 package software.amazon.smithy.model.validation.validators;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
+import java.util.Map;
+import java.util.Set;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.NeighborProviderIndex;
 import software.amazon.smithy.model.neighbor.NeighborProvider;
 import software.amazon.smithy.model.selector.Selector;
 import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.model.traits.TraitDefinition;
 import software.amazon.smithy.model.validation.AbstractValidator;
@@ -40,44 +45,69 @@
 public final class TraitTargetValidator extends AbstractValidator {
     @Override
     public List<ValidationEvent> validate(Model model) {
+        Collection<SelectorTest> tests = createTests(model);
+        List<ValidationEvent> events = new ArrayList<>();
         NeighborProvider neighborProvider = model.getKnowledge(NeighborProviderIndex.class).getProvider();
-        return model.shapes()
-                .flatMap(shape -> getSelectors(shape, model))
-                .filter(check -> !matchesSelector(check, model, neighborProvider))
-                .map(check -> error(check.shape, String.format(
-                        "Trait `%s` cannot be applied to `%s`. This trait may only be applied to shapes "
-                        + "that match the following selector: %s",
-                        Trait.getIdiomaticTraitName(check.trait.toShapeId()),
-                        check.shape.getId(), check.selector)))
-                .collect(Collectors.toList());
-    }
 
-    private static final class SelectorCheck {
-        final Shape shape;
-        final Trait trait;
-        final Selector selector;
+        for (SelectorTest test : tests) {
+            // Find the shapes that this trait can be applied to.
+            Set<Shape> matches = test.selector.select(model, neighborProvider);
 
-        SelectorCheck(Shape shape, Trait trait, Selector selector) {
-            this.shape = shape;
-            this.trait = trait;
-            this.selector = selector;
+            // Remove the allowed locations from the real locations, leaving only
+            // the shapes in the set that are invalid.
+            test.appliedTo.removeAll(matches);
+
+            for (Shape shape : test.appliedTo) {
+                events.add(error(shape, test.trait, String.format(
+                        "Trait `%s` cannot be applied to `%s`. This trait may only be applied "
+                        + "to shapes that match the following selector: %s",
+                        Trait.getIdiomaticTraitName(test.trait.toShapeId()),
+                        shape.getId(), test.selector)));
+            }
         }
+
+        return events;
     }
 
-    private Stream<SelectorCheck> getSelectors(Shape shape, Model model) {
-        return shape.getAllTraits().values().stream()
-                .map(trait -> new SelectorCheck(shape, trait, resolveSelector(trait, model)));
+    private Collection<SelectorTest> createTests(Model model) {
+        Map<ShapeId, SelectorTest> tests = new HashMap<>();
+        Map<ShapeId, Selector> selectors = new HashMap<>();
+
+        for (Shape shape : model.toSet()) {
+            for (Trait trait : shape.getAllTraits().values()) {
+                // The trait selector has to be resolved against the model,
+                // and possibly defaulted. This just caches that result since
+                // it's called for every single trait applied to a shape.
+                Selector selector = selectors.computeIfAbsent(trait.toShapeId(), id -> resolveSelector(id, model));
+
+                // Only need to test the location for traits that have some
+                // kind of constraint.
+                if (!selector.toString().trim().equals("*")) {
+                    SelectorTest test = tests.computeIfAbsent(
+                            trait.toShapeId(),
+                            id -> new SelectorTest(trait, selector));
+                    test.appliedTo.add(shape);
+                }
+            }
+        }
+
+        return tests.values();
     }
 
-    private Selector resolveSelector(Trait trait, Model model) {
-        return model.getTraitDefinition(trait).map(TraitDefinition::getSelector).orElse(Selector.IDENTITY);
+    private Selector resolveSelector(ShapeId id, Model model) {
+        return model.getTraitDefinition(id)
+                .map(TraitDefinition::getSelector)
+                .orElse(Selector.IDENTITY);
     }
 
-    private boolean matchesSelector(
-            SelectorCheck check,
-            Model model,
-            NeighborProvider neighborProvider
-    ) {
-        return check.selector.select(model, neighborProvider).contains(check.shape);
+    private static final class SelectorTest {
+        final Trait trait;
+        final Selector selector;
+        final Set<Shape> appliedTo = new HashSet<>();
+
+        SelectorTest(Trait trait, Selector selector) {
+            this.trait = trait;
+            this.selector = selector;
+        }
     }
 }