Generalize indexing of planner rules to support multiple concrete classes per rule

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/CascadesRule.java

@@ -27,7 +27,6 @@
 
 import javax.annotation.Nonnull;
 import java.util.Collection;
-import java.util.Optional;
 import java.util.Set;
 
 /**
@@ -76,16 +75,16 @@
         this.requirementDependencies = ImmutableSet.copyOf(requirementDependencies);
     }
 
     /**
-     * Returns the class of the operator at the root of the binding expression, if this rule uses a non-trivial binding.
-     * Used primarily for indexing rules for more efficient rule search.
-     * @return the class of the root of this rule's binding, or <code>Optional.empty()</code> if the rule matches anything
-     * @see PlanningRuleSet
+     * {@inheritDoc}
+     *
+     * <p>By default, this is derived from {@link BindingMatcher#getRootClasses()}. Subclasses that need to opt into the
+     * always-rules bucket should override this to return an empty set.
      */
     @Nonnull
     @Override
-    public Optional<Class<?>> getRootOperator() {
-        return Optional.of(matcher.getRootClass());
+    public Set<Class<?>> getRootOperators() {
+        return matcher.getRootClasses();
     }
 
     @Nonnull

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/PlannerRule.java

@@ -23,8 +23,8 @@
 import com.apple.foundationdb.annotation.API;
 import com.apple.foundationdb.record.query.plan.cascades.matching.structure.BindingMatcher;
 
 import javax.annotation.Nonnull;
-import java.util.Optional;
+import java.util.Set;
 
 /**
  * Basic rule interface.
@@ -37,13 +37,15 @@
 @API(API.Status.EXPERIMENTAL)
 public interface PlannerRule<C extends PlannerRuleCall, T> {
     /**
-     * Returns the class of the operator at the root of the binding expression, if this rule uses a non-trivial binding.
-     * Used primarily for indexing rules for more efficient rule search.
-     * @return the class of the root of this rule's binding, or <code>Optional.empty()</code> if the rule matches anything
+     * Returns the set of concrete operator classes that this rule should be indexed under in a rule set. This is used
+     * by the planner to quickly pick the subset of rules that could possibly fire on a visited node. If the set is
+     * empty, the rule will go into the always-rules bucket. Otherwise, the rule is indexed under each class returned.
+     * @return a (possibly empty) set of classes
      * @see PlanningRuleSet
+     * @see BindingMatcher#getRootClasses()
      */
     @Nonnull
-    Optional<Class<?>> getRootOperator();
+    Set<Class<?>> getRootOperators();
 
     void onMatch(@Nonnull C call);
 

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/matching/structure/BindingMatcher.java

@@ -29,7 +29,8 @@
 import javax.annotation.Nonnull;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
+import java.util.Set;
 import java.util.stream.Stream;
 
 /**
@@ -52,15 +53,27 @@
     String INDENTATION = "    ";
 
     /**
-     * Get a class that this matcher can match. Ideally, it should be the lowest such class, but it may not be.
-     * A planner will generally use this method to quickly determine a set of rules that could match an expression,
-     * without considering each rule and trying to apply it. A good implementation of this method helps the planner
-     * match rules efficiently.
-     * @return a class object for a class that is a super class of every planner expression this matcher can match
+     * Returns the dispatch class for this matcher, i.e., the upper bound of the type {@code T} that this matcher binds
+     * to. The class returned here is used by {@link #bindMatches} to gate the dispatching; it is also the class on
+     * which {@code where(…)} constraints operate. For purposes of rule-set indexing (i.e., picking the subset of rules
+     * that could possibly fire on a given visited node), see {@link #getRootClasses()}, which can be narrower than the
+     * upper-bound class returned here.
+     * @return the class object for the dispatch class
      */
     @Nonnull
     Class<T> getRootClass();
 
+    /**
+     * Get the set of concrete classes under which the enclosing rule of this matcher should be indexed in a rule set.
+     *
+     * <p>The default returns {@code Set.of(getRootClass())}, which is appropriate for most matchers. Matchers that need
+     * to bind more than one subclass should override this in order to benefit from targeted indexing.
+     */
+    @Nonnull
+    default Set<Class<?>> getRootClasses() {
+        return Set.of(getRootClass());
+    }
+
     /**
      * Attempt to match this matcher against the given object.
      * Note that implementations should only attempt to match the given object with this matcher and delegate
@@ -105,7 +118,8 @@
     default BindingMatcher<T> where(@Nonnull final BindingMatcher<? super T> downstream) {
         return TypedMatcherWithExtractAndDownstream.typedWithDownstream(getRootClass(),
                 Extractor.identity(),
-                AllOfMatcher.matchingAllOf(getRootClass(), ImmutableList.of(this, downstream)));
+                AllOfMatcher.matchingAllOf(getRootClass(), ImmutableList.of(this, downstream)),
+                getRootClasses());
     }
 
     /**

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/matching/structure/TypedMatcher.java

@@ -21,9 +21,14 @@
 package com.apple.foundationdb.record.query.plan.cascades.matching.structure;
 
 import com.apple.foundationdb.annotation.API;
 import com.apple.foundationdb.record.query.plan.RecordQueryPlannerConfiguration;
+import com.apple.foundationdb.record.query.plan.cascades.debug.Debugger;
+import com.google.common.base.Verify;
+import com.google.common.collect.ImmutableSet;
 
 import javax.annotation.Nonnull;
+import java.util.Set;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 /**
@@ -40,31 +45,74 @@
 public class TypedMatcher<T> implements BindingMatcher<T> {
     @Nonnull
     private final Class<T> bindableClass;
+    @Nonnull
+    private final Set<Class<?>> rootClasses;
 
     public TypedMatcher(@Nonnull final Class<T> bindableClass) {
+        this(bindableClass, ImmutableSet.of(bindableClass));
+    }
+
+    public TypedMatcher(@Nonnull final Class<T> bindableClass,
+                        @Nonnull final Set<? extends Class<?>> rootClasses) {
+        // Sanity-check that every advertised root class is a (sub-)type of `bindableClass`.
+        Debugger.sanityCheck(() -> rootClasses.forEach(c -> Verify.verify(bindableClass.isAssignableFrom(c),
+                "rootClass %s is not assignable from bindableClass %s", c, bindableClass)));
+
         this.bindableClass = bindableClass;
+        this.rootClasses = ImmutableSet.copyOf(rootClasses);
     }
 
     @Nonnull
     @Override
-    public Class<T> getRootClass() {
+    public final Class<T> getRootClass() {
         return bindableClass;
     }
 
+    @Nonnull
+    @Override
+    public final Set<Class<?>> getRootClasses() {
+        return rootClasses;
+    }
+
+    /**
+     * Returns whether the {@link #rootClasses} set contains only {@link #bindableClass} itself.
+     */
+    private boolean rootClassesIsTrivial() {
+        return rootClasses.size() == 1 && rootClasses.contains(bindableClass);
+    }
+
     @Nonnull
     @Override
     public Stream<PlannerBindings> bindMatchesSafely(@Nonnull RecordQueryPlannerConfiguration plannerConfiguration, @Nonnull PlannerBindings outerBindings, @Nonnull T in) {
+        // If `rootClasses` is not the trivial default `{bindableClass}`, sanity-check that the class of `in` is
+        // (exactly) one of the `rootClasses`. When `rootClasses` is a narrower set than the default, the matcher
+        // relies on the rule-set indexing path to only invoke it for instances of one of those classes.
+        //
+        // In the trivial default case this check is skipped on purpose, as `bindableClass` may be an abstract base
+        // class, and the check would then fail if `in.getClass()` is a concrete subclass.
+        if (!rootClassesIsTrivial()) {
+            Debugger.sanityCheck(() -> Verify.verify(rootClasses.contains(in.getClass()),
+                    "`TypedMatcher` invoked with class `%s`; expected one of %s", in.getClass(), rootClasses));
+        }
+
         return Stream.of(PlannerBindings.from(this, in));
     }
 
     @Override
     public String explainMatcher(@Nonnull final Class<?> atLeastType, @Nonnull final String boundId, @Nonnull final String indentation) {
-        if (getRootClass().isAssignableFrom(atLeastType)) {
+        // Don't bother emitting a redundant `case _: «RootClass»` type guard if it’s already established (according to
+        // `atLeastType`) that `boundId` is at least of type `getRootClass()`. Just emit a `_` wildcard pattern.
+        if (rootClassesIsTrivial() && getRootClass().isAssignableFrom(atLeastType)) {
             return "case _ => success ";
-        } else {
-            return "case _: " + getRootClass().getSimpleName() + " => success ";
         }
+
+        // Emit a typed wildcard pattern. Print a `(… | …)` union type if there are multiple root classes.
+        final String alternatives = rootClasses.stream()
+                .map(Class::getSimpleName)
+                .collect(Collectors.joining(" | "));
+        final String pattern = rootClasses.size() > 1 ? "(" + alternatives + ")" : alternatives;
+        return "case _: " + pattern + " => success ";
     }
 
     @Nonnull
     public static <T> TypedMatcher<T> typed(@Nonnull final Class<T> bindableClass) {

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/matching/structure/TypedMatcherWithExtractAndDownstream.java

@@ -23,7 +23,8 @@
 import com.apple.foundationdb.annotation.API;
 import com.apple.foundationdb.record.query.plan.RecordQueryPlannerConfiguration;
 
 import javax.annotation.Nonnull;
+import java.util.Set;
 import java.util.stream.Stream;
 
 import static com.apple.foundationdb.record.query.plan.cascades.matching.structure.BindingMatcher.newLine;
@@ -48,6 +49,15 @@
         this.downstream = downstream;
     }
 
+    protected TypedMatcherWithExtractAndDownstream(@Nonnull final Class<T> bindableClass,
+                                                   @Nonnull final Extractor<? super T, ?> extractor,
+                                                   @Nonnull final BindingMatcher<?> downstream,
+                                                   @Nonnull final Set<? extends Class<?>> rootClasses) {
+        super(bindableClass, rootClasses);
+        this.extractor = extractor;
+        this.downstream = downstream;
+    }
+
     @Nonnull
     public Extractor<? super T, ?> getExtractor() {
         return extractor;
@@ -91,4 +101,12 @@
                                                                                                @Nonnull final BindingMatcher<?> downstream) {
         return new TypedMatcherWithExtractAndDownstream<>(bindableClass, extractor, downstream);
     }
+
+    @Nonnull
+    public static <S, T extends S> TypedMatcherWithExtractAndDownstream<T> typedWithDownstream(@Nonnull final Class<T> bindableClass,
+                                                                                               @Nonnull final Extractor<? super T, ?> extractor,
+                                                                                               @Nonnull final BindingMatcher<?> downstream,
+                                                                                               @Nonnull final Set<? extends Class<?>> concreteRootClasses) {
+        return new TypedMatcherWithExtractAndDownstream<>(bindableClass, extractor, downstream, concreteRootClasses);
+    }
 }

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/matching/structure/ValueMatchers.java

@@ -36,7 +36,8 @@
 import com.apple.foundationdb.record.query.plan.cascades.values.Value;
 import com.apple.foundationdb.record.query.plan.cascades.values.VariadicFunctionValue;
 import com.apple.foundationdb.tuple.TupleOrdering;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 
 import javax.annotation.Nonnull;
 import java.util.Arrays;
@@ -57,11 +58,24 @@
         // do not instantiate
     }
 
+    /**
+     * Returns a matcher that binds any {@link Value}.
+     */
     @Nonnull
     public static BindingMatcher<Value> anyValue() {
         return typed(Value.class);
     }
 
+    /**
+     * Returns a matcher that binds any {@link Value} whose concrete class is one of the given {@code subclasses}.
+     * The resulting matcher advertises the {@code subclasses} via {@link BindingMatcher#getRootClasses()} so that the
+     * rule set indexes the enclosing rule under each concrete subclass individually.
+     */
+    @Nonnull
+    public static BindingMatcher<Value> anyValueOfType(@Nonnull final ImmutableSet<Class<? extends Value>> subclasses) {
+        return new TypedMatcher<>(Value.class, subclasses);
+    }
+
     @Nonnull
     public static BindingMatcher<FieldValue> anyFieldValue() {
         return typed(FieldValue.class);

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/predicates/simplification/AbsorptionRule.java

@@ -33,8 +33,8 @@
 import com.google.common.collect.Sets;
 
 import javax.annotation.Nonnull;
 import java.util.Collection;
-import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 import static com.apple.foundationdb.record.query.plan.cascades.matching.structure.MultiMatcher.all;
@@ -68,8 +68,8 @@
 
     @Nonnull
     @Override
-    public Optional<Class<?>> getRootOperator() {
-        return Optional.empty();
+    public Set<Class<?>> getRootOperators() {
+        return Set.of();
     }
 
     @Override

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/predicates/simplification/AnnulmentAndRule.java

@@ -28,7 +28,6 @@
 import com.apple.foundationdb.record.query.plan.cascades.predicates.QueryPredicate;
 
 import javax.annotation.Nonnull;
-import java.util.Optional;
 
 import static com.apple.foundationdb.record.query.plan.cascades.matching.structure.MultiMatcher.all;
 import static com.apple.foundationdb.record.query.plan.cascades.matching.structure.QueryPredicateMatchers.anyPredicate;
@@ -50,12 +49,6 @@ public AnnulmentAndRule() {
         super(rootMatcher);
     }
 
-    @Nonnull
-    @Override
-    public Optional<Class<?>> getRootOperator() {
-        return Optional.of(AndPredicate.class);
-    }
-
     @Override
     public void onMatch(@Nonnull final QueryPredicateSimplificationRuleCall call) {
         final var bindings = call.getBindings();

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/predicates/simplification/AnnulmentOrRule.java

@@ -28,7 +28,6 @@
 import com.apple.foundationdb.record.query.plan.cascades.predicates.QueryPredicate;
 
 import javax.annotation.Nonnull;
-import java.util.Optional;
 
 import static com.apple.foundationdb.record.query.plan.cascades.matching.structure.MultiMatcher.all;
 import static com.apple.foundationdb.record.query.plan.cascades.matching.structure.QueryPredicateMatchers.anyPredicate;
@@ -51,12 +50,6 @@ public AnnulmentOrRule() {
         super(rootMatcher);
     }
 
-    @Nonnull
-    @Override
-    public Optional<Class<?>> getRootOperator() {
-        return Optional.of(OrPredicate.class);
-    }
-
     @Override
     public void onMatch(@Nonnull final QueryPredicateSimplificationRuleCall call) {
         final var bindings = call.getBindings();

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/predicates/simplification/DeMorgansTheoremRule.java

@@ -33,7 +33,6 @@
 
 import javax.annotation.Nonnull;
 import java.util.Collection;
-import java.util.Optional;
 
 import static com.apple.foundationdb.record.query.plan.cascades.matching.structure.MultiMatcher.all;
 import static com.apple.foundationdb.record.query.plan.cascades.matching.structure.QueryPredicateMatchers.anyPredicate;
@@ -65,12 +64,6 @@ public DeMorgansTheoremRule(@Nonnull final Class<P> majorClass,
         this.termMatcher = termMatcher;
     }
 
-    @Nonnull
-    @Override
-    public Optional<Class<?>> getRootOperator() {
-        return Optional.of(NotPredicate.class);
-    }
-
     @Override
     public void onMatch(@Nonnull final QueryPredicateSimplificationRuleCall call) {
         final var bindings = call.getBindings();

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/predicates/simplification/IdentityAndRule.java

@@ -28,7 +28,6 @@
 import com.google.common.collect.ImmutableList;
 
 import javax.annotation.Nonnull;
-import java.util.Optional;
 
 import static com.apple.foundationdb.record.query.plan.cascades.matching.structure.MultiMatcher.all;
 import static com.apple.foundationdb.record.query.plan.cascades.matching.structure.QueryPredicateMatchers.anyPredicate;
@@ -50,12 +49,6 @@ public IdentityAndRule() {
         super(rootMatcher);
     }
 
-    @Nonnull
-    @Override
-    public Optional<Class<?>> getRootOperator() {
-        return Optional.of(AndPredicate.class);
-    }
-
     @Override
     public void onMatch(@Nonnull final QueryPredicateSimplificationRuleCall call) {
         final var bindings = call.getBindings();

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/predicates/simplification/IdentityOrRule.java

@@ -28,7 +28,6 @@
 import com.google.common.collect.ImmutableList;
 
 import javax.annotation.Nonnull;
-import java.util.Optional;
 
 import static com.apple.foundationdb.record.query.plan.cascades.matching.structure.MultiMatcher.all;
 import static com.apple.foundationdb.record.query.plan.cascades.matching.structure.QueryPredicateMatchers.anyPredicate;
@@ -50,12 +49,6 @@ public IdentityOrRule() {
         super(rootMatcher);
     }
 
-    @Nonnull
-    @Override
-    public Optional<Class<?>> getRootOperator() {
-        return Optional.of(OrPredicate.class);
-    }
-
     @Override
     public void onMatch(@Nonnull final QueryPredicateSimplificationRuleCall call) {
         final var bindings = call.getBindings();