perf(constraint-streams): fused equal indexer (#2346)

Author: triceo

core/src/main/java/ai/timefold/solver/core/impl/bavet/bi/joiner/BiJoinerComber.java

@@ -6,83 +6,87 @@
 
 import ai.timefold.solver.core.api.score.stream.bi.BiJoiner;
 
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
 /**
  * Combs an array of {@link BiJoiner} instances into a mergedJoiner and a mergedFiltering.
  *
  * @param <A>
  * @param <B>
  */
+@NullMarked
 public final class BiJoinerComber<A, B> {
 
+    @SafeVarargs
     public static <A, B> BiJoinerComber<A, B> comb(BiJoiner<A, B>... joiners) {
-        List<DefaultBiJoiner<A, B>> defaultJoinerList = new ArrayList<>(joiners.length);
-        List<BiPredicate<A, B>> filteringList = new ArrayList<>(joiners.length);
+        var defaultJoinerList = new ArrayList<DefaultBiJoiner<A, B>>(joiners.length);
+        var filteringList = new ArrayList<BiPredicate<A, B>>(joiners.length);
 
-        int indexOfFirstFilter = -1;
+        var indexOfFirstFilter = -1;
         // Make sure all indexing joiners, if any, come before filtering joiners. This is necessary for performance.
-        for (int i = 0; i < joiners.length; i++) {
-            BiJoiner<A, B> joiner = joiners[i];
+        for (var i = 0; i < joiners.length; i++) {
+            var joiner = joiners[i];
             if (joiner instanceof FilteringBiJoiner) {
                 // From now on, only allow filtering joiners.
                 indexOfFirstFilter = i;
                 filteringList.add(((FilteringBiJoiner<A, B>) joiner).getFilter());
             } else if (joiner instanceof DefaultBiJoiner) {
                 if (indexOfFirstFilter >= 0) {
-                    throw new IllegalStateException("Indexing joiner (" + joiner + ") must not follow " +
-                            "a filtering joiner (" + joiners[indexOfFirstFilter] + ").\n" +
-                            "Maybe reorder the joiners such that filtering() joiners are later in the parameter list.");
+                    throw new IllegalStateException("""
+                            Indexing joiner (%s) must not follow a filtering joiner (%s).
+                            Maybe reorder the joiners such that filtering() joiners are later in the parameter list."""
+                            .formatted(joiner, joiners[indexOfFirstFilter]));
                 }
                 defaultJoinerList.add((DefaultBiJoiner<A, B>) joiner);
             } else {
-                throw new IllegalArgumentException("The joiner class (" + joiner.getClass() + ") is not supported.");
+                throw new IllegalArgumentException("The joiner class (%s) is not supported."
+                        .formatted(joiner.getClass()));
             }
         }
-        DefaultBiJoiner<A, B> mergedJoiner = DefaultBiJoiner.merge(defaultJoinerList);
-        BiPredicate<A, B> mergedFiltering = mergeFiltering(filteringList);
+        var mergedJoiner = DefaultBiJoiner.merge(defaultJoinerList);
+        var mergedFiltering = mergeFiltering(filteringList);
         return new BiJoinerComber<>(mergedJoiner, mergedFiltering);
     }
 
-    private static <A, B> BiPredicate<A, B> mergeFiltering(List<BiPredicate<A, B>> filteringList) {
-        if (filteringList.isEmpty()) {
-            return null;
-        }
-        switch (filteringList.size()) {
-            case 1:
-                return filteringList.get(0);
-            case 2:
-                return filteringList.get(0).and(filteringList.get(1));
-            default:
-                // Avoid predicate.and() when more than 2 predicates for debugging and potentially performance
-                return (A a, B b) -> {
-                    for (BiPredicate<A, B> predicate : filteringList) {
+    @SuppressWarnings("unchecked")
+    private static <A, B> @Nullable BiPredicate<A, B> mergeFiltering(List<BiPredicate<A, B>> filteringList) {
+        return switch (filteringList.size()) {
+            case 0 -> null;
+            case 1 -> filteringList.getFirst();
+            default -> {
+                // Avoid predicate.and() for debugging and potential performance
+                var filteringArray = filteringList.toArray(new BiPredicate[0]);
+                yield (A a, B b) -> {
+                    for (var predicate : filteringArray) {
                         if (!predicate.test(a, b)) {
                             return false;
                         }
                     }
                     return true;
                 };
-        }
+            }
+        };
     }
 
     private DefaultBiJoiner<A, B> mergedJoiner;
-    private final BiPredicate<A, B> mergedFiltering;
+    private final @Nullable BiPredicate<A, B> mergedFiltering;
 
-    public BiJoinerComber(DefaultBiJoiner<A, B> mergedJoiner, BiPredicate<A, B> mergedFiltering) {
+    public BiJoinerComber(DefaultBiJoiner<A, B> mergedJoiner, @Nullable BiPredicate<A, B> mergedFiltering) {
         this.mergedJoiner = mergedJoiner;
         this.mergedFiltering = mergedFiltering;
     }
 
     /**
-     * @return never null
+     * Returns the merged indexing joiner,
+     * reordered equal-first so the indexer chain always has its (merged) equal level at the top.
+     * Computed on read to also cover {@link #addJoiner} appends.
      */
     public DefaultBiJoiner<A, B> getMergedJoiner() {
-        return mergedJoiner;
+        return mergedJoiner.reorderedEqualsFirst();
     }
 
-    /**
-     * @return null if not applicable
-     */
-    public BiPredicate<A, B> getMergedFiltering() {
+    public @Nullable BiPredicate<A, B> getMergedFiltering() {
         return mergedFiltering;
     }
 

core/src/main/java/ai/timefold/solver/core/impl/bavet/bi/joiner/DefaultBiJoiner.java

@@ -9,9 +9,12 @@
 import ai.timefold.solver.core.impl.bavet.common.joiner.AbstractJoiner;
 import ai.timefold.solver.core.impl.bavet.common.joiner.JoinerType;
 
-import org.jspecify.annotations.NonNull;
+import org.jspecify.annotations.NullMarked;
 
-public final class DefaultBiJoiner<A, B> extends AbstractJoiner<B> implements BiJoiner<A, B> {
+@NullMarked
+public final class DefaultBiJoiner<A, B>
+        extends AbstractJoiner<B>
+        implements BiJoiner<A, B> {
 
     private static final DefaultBiJoiner NONE = new DefaultBiJoiner(new Function[0], new JoinerType[0], new Function[0]);
 
@@ -28,40 +31,64 @@ public DefaultBiJoiner(Function<A, ?>[] leftMappings, JoinerType[] joinerTypes,
     }
 
     public static <A, B> DefaultBiJoiner<A, B> merge(List<DefaultBiJoiner<A, B>> joinerList) {
-        if (joinerList.size() == 1) {
-            return joinerList.get(0);
-        }
-        return joinerList.stream().reduce(NONE, DefaultBiJoiner::and);
+        return switch (joinerList.size()) {
+            case 0 -> NONE;
+            case 1 -> joinerList.getFirst();
+            default -> joinerList.stream().reduce(NONE, DefaultBiJoiner::and);
+        };
     }
 
     @Override
-    public @NonNull DefaultBiJoiner<A, B> and(@NonNull BiJoiner<A, B> otherJoiner) {
-        DefaultBiJoiner<A, B> castJoiner = (DefaultBiJoiner<A, B>) otherJoiner;
-        int joinerCount = getJoinerCount();
-        int castJoinerCount = castJoiner.getJoinerCount();
-        int newJoinerCount = joinerCount + castJoinerCount;
-        JoinerType[] newJoinerTypes = Arrays.copyOf(this.joinerTypes, newJoinerCount);
-        Function[] newLeftMappings = Arrays.copyOf(this.leftMappings, newJoinerCount);
-        Function[] newRightMappings = Arrays.copyOf(this.rightMappings, newJoinerCount);
-        for (int i = 0; i < castJoinerCount; i++) {
-            int newJoinerIndex = i + joinerCount;
+    public DefaultBiJoiner<A, B> and(BiJoiner<A, B> otherJoiner) {
+        var castJoiner = (DefaultBiJoiner<A, B>) otherJoiner;
+        var joinerCount = getJoinerCount();
+        var castJoinerCount = castJoiner.getJoinerCount();
+        var newJoinerCount = joinerCount + castJoinerCount;
+        var newJoinerTypes = Arrays.copyOf(this.joinerTypes, newJoinerCount);
+        var newLeftMappings = Arrays.copyOf(this.leftMappings, newJoinerCount);
+        var newRightMappings = Arrays.copyOf(this.rightMappings, newJoinerCount);
+        for (var i = 0; i < castJoinerCount; i++) {
+            var newJoinerIndex = i + joinerCount;
             newJoinerTypes[newJoinerIndex] = castJoiner.getJoinerType(i);
             newLeftMappings[newJoinerIndex] = castJoiner.getLeftMapping(i);
             newRightMappings[newJoinerIndex] = castJoiner.getRightMapping(i);
         }
         return new DefaultBiJoiner<>(newLeftMappings, newJoinerTypes, newRightMappings);
     }
 
+    /**
+     * @return this if already equal-first (or single joiner); otherwise a copy with all
+     *         {@link JoinerType#EQUAL} joiners moved to the front (stable, see
+     *         {@link AbstractJoiner#equalsFirstSortedPositions}).
+     */
+    public DefaultBiJoiner<A, B> reorderedEqualsFirst() {
+        var order = equalsFirstSortedPositions(joinerTypes);
+        if (order == null) {
+            return this;
+        }
+        var count = order.length;
+        var newLeftMappings = new Function[count];
+        var newJoinerTypes = new JoinerType[count];
+        var newRightMappings = new Function[count];
+        for (var i = 0; i < count; i++) {
+            var from = order[i];
+            newLeftMappings[i] = leftMappings[from];
+            newJoinerTypes[i] = joinerTypes[from];
+            newRightMappings[i] = rightMappings[from];
+        }
+        return new DefaultBiJoiner<>(newLeftMappings, newJoinerTypes, newRightMappings);
+    }
+
     public Function<A, Object> getLeftMapping(int index) {
         return leftMappings[index];
     }
 
     public boolean matches(A a, B b) {
-        int joinerCount = getJoinerCount();
-        for (int i = 0; i < joinerCount; i++) {
-            JoinerType joinerType = getJoinerType(i);
-            Object leftMapping = getLeftMapping(i).apply(a);
-            Object rightMapping = getRightMapping(i).apply(b);
+        var joinerCount = getJoinerCount();
+        for (var i = 0; i < joinerCount; i++) {
+            var joinerType = getJoinerType(i);
+            var leftMapping = getLeftMapping(i).apply(a);
+            var rightMapping = getRightMapping(i).apply(b);
             if (!joinerType.matches(leftMapping, rightMapping)) {
                 return false;
             }