TimefoldAI · ge0ffrey · Dec 26, 2025 · Dec 26, 2025 · Dec 26, 2025 · Dec 27, 2025
@@ -1,5 +1,6 @@
 package ai.timefold.solver.core.api.score.stream;
 
+import java.util.Collection;
 import java.util.function.BiFunction;
 import java.util.function.BiPredicate;
 import java.util.function.Function;
@@ -35,7 +36,6 @@
 public final class Joiners {
 
     // TODO Support using non-natural comparators, such as lessThan(leftMapping, rightMapping, comparator).
-    // TODO Support collection-based joiners, such as containing(), intersecting() and disjoint().
 
     // ************************************************************************
     // BiJoiner
@@ -78,6 +78,24 @@ public final class Joiners {
         return new DefaultBiJoiner<>(leftMapping, JoinerType.EQUAL, rightMapping);
     }
 
+    // TODO javadoc
+    public static <A, B, Property_> @NonNull BiJoiner<A, B> contain(Function<A, Collection<Property_>> leftMapping,
+            Function<B, Property_> rightMapping) {
+        return new DefaultBiJoiner<>(leftMapping, JoinerType.CONTAIN, rightMapping);
+    }
+
+    // TODO javadoc
+    public static <A, B, Property_> @NonNull BiJoiner<A, B> containedIn(Function<A, Property_> leftMapping,
+            Function<B, Collection<Property_>> rightMapping) {
+        return new DefaultBiJoiner<>(leftMapping, JoinerType.CONTAINED_IN, rightMapping);
+    }
+
+    // TODO javadoc
+    public static <A, B, Property_> @NonNull BiJoiner<A, B> containAny(Function<A, Collection<Property_>> leftMapping,
+            Function<B, Collection<Property_>> rightMapping) {
+        return new DefaultBiJoiner<>(leftMapping, JoinerType.CONTAIN_ANY, rightMapping);
+    }
+
     /**
      * As defined by {@link #lessThan(Function, Function)} with both arguments using the same mapping.
      *

@@ -15,18 +15,16 @@ public final class DefaultBiJoiner<A, B> extends AbstractJoiner<B> implements Bi
 
     private static final DefaultBiJoiner NONE = new DefaultBiJoiner(new Function[0], new JoinerType[0], new Function[0]);
 
-    private final Function<A, ?>[] leftMappings;
+    private final Function<A, Object>[] leftMappings;
 
-    public <Property_> DefaultBiJoiner(Function<A, Property_> leftMapping, JoinerType joinerType,
-            Function<B, Property_> rightMapping) {
+    public DefaultBiJoiner(Function<A, ?> leftMapping, JoinerType joinerType, Function<B, ?> rightMapping) {
         super(rightMapping, joinerType);
         this.leftMappings = new Function[] { leftMapping };
     }
 
-    public <Property_> DefaultBiJoiner(Function<A, Property_>[] leftMappings, JoinerType[] joinerTypes,
-            Function<B, Property_>[] rightMappings) {
+    public DefaultBiJoiner(Function<A, ?>[] leftMappings, JoinerType[] joinerTypes, Function<B, ?>[] rightMappings) {
         super(rightMappings, joinerTypes);
-        this.leftMappings = leftMappings;
+        this.leftMappings = (Function<A, Object>[]) Objects.requireNonNull(leftMappings);
     }
 
     public static <A, B> DefaultBiJoiner<A, B> merge(List<DefaultBiJoiner<A, B>> joinerList) {
@@ -55,7 +53,7 @@ public static <A, B> DefaultBiJoiner<A, B> merge(List<DefaultBiJoiner<A, B>> joi
     }
 
     public Function<A, Object> getLeftMapping(int index) {
-        return (Function<A, Object>) leftMappings[index];
+        return leftMappings[index];
     }
 
     public boolean matches(A a, B b) {

@@ -25,31 +25,13 @@ final class ComparisonIndexer<T, Key_ extends Comparable<Key_>>
     private final NavigableMap<Key_, Indexer<T>> comparisonMap;
 
     /**
-     * Construct an {@link ComparisonIndexer} which immediately ends in the backend.
-     * This means {@code compositeKey} must be a single key.
-     *
      * @param comparisonJoinerType the type of comparison to use
+     * @param keyRetriever determines if it immediately goes to a {@link IndexerBackend} or if it uses a {@link CompositeKey}.
      * @param downstreamIndexerSupplier the supplier of the downstream indexer
      */
-    public ComparisonIndexer(JoinerType comparisonJoinerType, Supplier<Indexer<T>> downstreamIndexerSupplier) {
-        this(comparisonJoinerType, new SingleKeyRetriever<>(), downstreamIndexerSupplier);
-    }
-
-    /**
-     * Construct an {@link ComparisonIndexer} which does not immediately go to a {@link IndexerBackend}.
-     * This means {@code compositeKey} must be an instance of {@link CompositeKey}.
-     *
-     * @param comparisonJoinerType the type of comparison to use
-     * @param keyIndex the index of the key to use within {@link CompositeKey}.
-     * @param downstreamIndexerSupplier the supplier of the downstream indexer
-     */
-    public ComparisonIndexer(JoinerType comparisonJoinerType, int keyIndex, Supplier<Indexer<T>> downstreamIndexerSupplier) {
-        this(comparisonJoinerType, new CompositeKeyRetriever<>(keyIndex), downstreamIndexerSupplier);
-    }
-
-    private ComparisonIndexer(JoinerType comparisonJoinerType, KeyRetriever<Key_> keyRetriever,
+    public ComparisonIndexer(JoinerType comparisonJoinerType, KeyRetriever<?> keyRetriever,
             Supplier<Indexer<T>> downstreamIndexerSupplier) {
-        this.keyRetriever = Objects.requireNonNull(keyRetriever);
+        this.keyRetriever = Objects.requireNonNull((KeyRetriever<Key_>) keyRetriever);
         this.downstreamIndexerSupplier = Objects.requireNonNull(downstreamIndexerSupplier);
         // For GT/GTE, the iteration order is reversed.
         // This allows us to iterate over the entire map from the start, stopping when the threshold is reached.

@@ -20,7 +20,7 @@
  * TriCompositeKey and higher are rare enough for {@link MegaCompositeKey} to suffice.
  */
 public sealed interface CompositeKey
-        permits MegaCompositeKey, BiCompositeKey {
+        permits BiCompositeKey, MegaCompositeKey {
 
     static CompositeKey none() {
         return MegaCompositeKey.EMPTY;
@@ -53,7 +53,7 @@ static CompositeKey ofMany(Object... keys) {
      * @param id Maps to a single {@link Indexer} instance in the indexer chain.
      * @return May be null if the key is null.
      * @param <Key_> {@link ComparisonIndexer} will expect this to implement {@link Comparable}.
-     *        {@link EqualsIndexer} will treat items as the same if they are equal.
+     *        {@link EqualIndexer} will treat items as the same if they are equal.
      */
     <Key_> Key_ get(int id);
 

@@ -0,0 +1,169 @@
+package ai.timefold.solver.core.impl.bavet.common.index;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+
+import ai.timefold.solver.core.impl.util.CompositeListEntry;
+import ai.timefold.solver.core.impl.util.ListEntry;
+import ai.timefold.solver.core.impl.util.Pair;
+
+import org.jspecify.annotations.NullMarked;
+
+@NullMarked
+final class ContainAnyIndexer<T, Key_, KeyCollection_ extends Collection<Key_>> implements Indexer<T> {
+
+    private final KeyRetriever<KeyCollection_> modifyKeyRetriever;
+    private final KeyRetriever<KeyCollection_> queryKeyRetriever;
+    private final Supplier<Indexer<T>> downstreamIndexerSupplier;
+    private final Map<Key_, Indexer<T>> downstreamIndexerMap = new HashMap<>();
+
+    /**
+     * @param keyRetriever determines if it immediately goes to a {@link IndexerBackend} or if it uses a {@link CompositeKey}.
+     * @param downstreamIndexerSupplier the supplier of the downstream indexer
+     */
+    public ContainAnyIndexer(KeyRetriever<Key_> keyRetriever, Supplier<Indexer<T>> downstreamIndexerSupplier) {
+        this.modifyKeyRetriever = Objects.requireNonNull((KeyRetriever<KeyCollection_>) keyRetriever);
+        this.queryKeyRetriever = Objects.requireNonNull((KeyRetriever<KeyCollection_>) keyRetriever);
+        this.downstreamIndexerSupplier = Objects.requireNonNull(downstreamIndexerSupplier);
+    }
+
+    @Override
+    public ListEntry<T> put(Object modifyCompositeKey, T tuple) {
+        KeyCollection_ indexKeyCollection = modifyKeyRetriever.apply(modifyCompositeKey);
+        List<Pair<Key_, ListEntry<T>>> children = new ArrayList<>(indexKeyCollection.size());
+        for (Key_ indexKey : indexKeyCollection) {
+            // Avoids computeIfAbsent in order to not create lambdas on the hot path.
+            Indexer<T> downstreamIndexer = downstreamIndexerMap.get(indexKey);
+            if (downstreamIndexer == null) {
+                downstreamIndexer = downstreamIndexerSupplier.get();
+                downstreamIndexerMap.put(indexKey, downstreamIndexer);
+            }
+            // Even though this method puts a tuple in multiple downstreamIndexers, it does not break size() or forEach()
+            // because at most one of those downstreamIndexers matches for a particular compositeKey
+            ListEntry<T> childListEntry = downstreamIndexer.put(modifyCompositeKey, tuple);
+            children.add(new Pair<>(indexKey, childListEntry));
+        }
+        return new CompositeListEntry<>(tuple, children);
+    }
+
+    @Override
+    public void remove(Object modifyCompositeKey, ListEntry<T> entry) {
+        KeyCollection_ indexKeyCollection = modifyKeyRetriever.apply(modifyCompositeKey);
+        List<Pair<Key_, ListEntry<T>>> children = ((CompositeListEntry<Key_, T>) entry).getChildren();
+        if (indexKeyCollection.size() != children.size()) {
+            throw new IllegalStateException(
+                    ("Impossible state: the tuple (%s) with composite key (%s) has a different number of children (%d)" +
+                            " than the index key collection size (%d).")
+                            .formatted(entry, modifyCompositeKey, children.size(), indexKeyCollection.size()));
+        }
+        for (Pair<Key_, ListEntry<T>> child : children) {
+            Key_ indexKey = child.key();
+            ListEntry<T> childListEntry = child.value();
+            // Avoids removeIfAbsent in order to not create lambdas on the hot path.
+            Indexer<T> downstreamIndexer = getDownstreamIndexer(modifyCompositeKey, indexKey);
+            downstreamIndexer.remove(modifyCompositeKey, childListEntry);
+            if (downstreamIndexer.isEmpty()) {
+                downstreamIndexerMap.remove(indexKey);
+            }
+        }
+    }
+
+    private Indexer<T> getDownstreamIndexer(Object compositeKey, Key_ indexerKey) {
+        Indexer<T> downstreamIndexer = downstreamIndexerMap.get(indexerKey);
+        if (downstreamIndexer == null) {
+            throw new IllegalStateException(
+                    "Impossible state: the composite key (%s) doesn't exist in the indexer %s."
+                            .formatted(compositeKey, this));
+        }
+        return downstreamIndexer;
+    }
+
+    @Override
+    public int size(Object queryCompositeKey) {
+        KeyCollection_ indexKeyCollection = queryKeyRetriever.apply(queryCompositeKey);
+        if (indexKeyCollection.isEmpty()) {
+            return 0;
+        } else if (indexKeyCollection.size() == 1) {
+            Indexer<T> downstreamIndexer = downstreamIndexerMap.get(indexKeyCollection.iterator().next());
+            return (downstreamIndexer == null) ? 0 : downstreamIndexer.size(queryCompositeKey);
+        } else {
+            AtomicInteger size = new AtomicInteger(0);
+            forEach(queryCompositeKey, tuple -> size.incrementAndGet());
+            return size.get();
+        }
+    }
+
+    @Override
+    public void forEach(Object queryCompositeKey, Consumer<T> tupleConsumer) {
+        KeyCollection_ indexKeyCollection = queryKeyRetriever.apply(queryCompositeKey);
+        if (indexKeyCollection.isEmpty()) {
+            return;
+        } else if (indexKeyCollection.size() == 1) {
+            Indexer<T> downstreamIndexer = downstreamIndexerMap.get(indexKeyCollection.iterator().next());
+            if (downstreamIndexer != null) {
+                downstreamIndexer.forEach(queryCompositeKey, tupleConsumer);
+            }
+        } else {
+            Set<T> distinctingSet = new HashSet<>(indexKeyCollection.size() * 16);
+            for (Key_ indexKey : indexKeyCollection) {
+                Indexer<T> downstreamIndexer = downstreamIndexerMap.get(indexKey);
+                if (downstreamIndexer != null) {
+                    downstreamIndexer.forEach(queryCompositeKey, tuple -> {
+                        if (distinctingSet.add(tuple)) {
+                            tupleConsumer.accept(tuple);
+                        }
+                    });
+                }
+            }
+        }
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return downstreamIndexerMap.isEmpty();
+    }
+
+    @Override
+    public List<? extends ListEntry<T>> asList(Object queryCompositeKey) {
+        KeyCollection_ indexKeyCollection = queryKeyRetriever.apply(queryCompositeKey);
+        if (indexKeyCollection.isEmpty()) {
+            return List.of();
+        } else if (indexKeyCollection.size() == 1) {
+            Key_ indexKey = indexKeyCollection.iterator().next();
+            Indexer<T> downstreamIndexer = downstreamIndexerMap.get(indexKey);
+            if (downstreamIndexer == null) {
+                return List.of();
+            }
+            return downstreamIndexer.asList(queryCompositeKey);
+        } else {
+            List<ListEntry<T>> list = new ArrayList<>(downstreamIndexerMap.size() * 16);
+            Set<T> distinctingSet = new HashSet<>(indexKeyCollection.size() * 16);
+            for (Key_ indexKey : indexKeyCollection) {
+                Indexer<T> downstreamIndexer = downstreamIndexerMap.get(indexKey);
+                if (downstreamIndexer != null) {
+                    downstreamIndexer.forEach(queryCompositeKey, tuple -> {
+                        if (distinctingSet.add(tuple)) {
+                            list.addAll(downstreamIndexer.asList(queryCompositeKey));
+                        }
+                    });
+                }
+            }
+            return list;
+        }
+    }
+
+    @Override
+    public String toString() {
+        return "size = " + downstreamIndexerMap.size();
+    }
+
+}