feat: introduce flatten() as a more generic alternative to flattenLast()

core/src/main/java/ai/timefold/solver/core/api/score/stream/bi/BiConstraintStream.java

@@ -7,6 +7,7 @@
 import static ai.timefold.solver.core.impl.util.ConstantLambdaUtils.uniConstantNull;
 
 import java.math.BigDecimal;
+import java.util.Collection;
 import java.util.Objects;
 import java.util.function.BiFunction;
 import java.util.function.BiPredicate;
@@ -1076,33 +1077,49 @@ public interface BiConstraintStream<A, B> extends ConstraintStream {
             @NonNull BiFunction<A, B, ResultD_> mappingD);
 
     /**
-     * Takes each tuple and applies a mapping on the last fact, which turns it into {@link Iterable}.
-     * Returns a constraint stream consisting of tuples of the first fact
-     * and the contents of the {@link Iterable} one after another.
+     * Takes each tuple and applies a mapping on its facts, which turns it into {@link Iterable}.
+     * Returns a constraint stream consisting of new tuples,
+     * each made of the original facts and one item from that iterable.
      * In other words, it will replace the current tuple with new tuples,
-     * a cartesian product of A and the individual items from the {@link Iterable}.
+     * a Cartesian product of (A, B) and the individual items from the {@link Iterable}.
      *
      * <p>
      * This may produce a stream with duplicate tuples.
      * See {@link #distinct()} for details.
      *
      * <p>
-     * In cases where the last fact is already {@link Iterable}, use {@link Function#identity()} as the argument.
-     *
-     * <p>
      * Simple example: assuming a constraint stream of {@code (PersonName, Person)}
      * {@code [(Ann, (name = Ann, roles = [USER, ADMIN])), (Beth, (name = Beth, roles = [USER])),
      * (Cathy, (name = Cathy, roles = [ADMIN, AUDITOR]))]},
-     * calling {@code flattenLast(Person::getRoles)} on such stream will produce a stream of
-     * {@code [(Ann, USER), (Ann, ADMIN), (Beth, USER), (Cathy, ADMIN), (Cathy, AUDITOR)]}.
-     *
-     * @param mapping function to convert the last fact in the original tuple into {@link Iterable}.
-     *        For performance, returning an implementation of {@link java.util.Collection} is preferred.
-     * @param <ResultB_> the type of the last fact in the resulting tuples.
+     * calling {@code flatten((name, person) -> person.getRoles()))} on such stream will produce a stream of
+     * {@code [(Ann, (name = Ann, roles = [USER, ADMIN]), USER),
+     * (Ann, (name = Ann, roles = [USER, ADMIN]), ADMIN),
+     * (Beth, (name = Beth, roles = [USER]), USER),
+     * (Cathy, (name = Cathy, roles = [ADMIN, AUDITOR]), ADMIN),
+     * (Cathy, (name = Cathy, roles = [ADMIN, AUDITOR]), AUDITOR)]}.
+     *
+     * @param mapping function to convert the original tuple into {@link Iterable}.
+     *        For performance, returning an implementation of {@link Collection} is preferred.
+     * @param <ResultC_> the type of the last fact in the resulting tuples.
      *        It is recommended that this type be deeply immutable.
      *        Not following this recommendation may lead to hard-to-debug hashing issues down the stream,
      *        especially if this value is ever used as a group key.
      */
+    <ResultC_> @NonNull TriConstraintStream<A, B, ResultC_>
+            flatten(@NonNull BiFunction<A, B, @NonNull Iterable<ResultC_>> mapping);
+
+    /**
+     * As defined by {@link #flatten(BiFunction)},
+     * only replacing the last fact in the original tuple by an item from the iterable.
+     * This means the resulting stream will still be a {@link BiConstraintStream},
+     * not a {@link TriConstraintStream}.
+     * <p>
+     * Simple example: assuming a constraint stream of {@code (PersonName, Person)}
+     * {@code [(Ann, (name = Ann, roles = [USER, ADMIN])), (Beth, (name = Beth, roles = [USER])),
+     * (Cathy, (name = Cathy, roles = [ADMIN, AUDITOR]))]},
+     * calling {@code flattenLast(Person::getRoles)} on such stream will produce a stream of
+     * {@code [(Ann, USER), (Ann, ADMIN), (Beth, USER), (Cathy, ADMIN), (Cathy, AUDITOR)]}.
+     */
     <ResultB_> @NonNull BiConstraintStream<A, ResultB_> flattenLast(@NonNull Function<B, @NonNull Iterable<ResultB_>> mapping);
 
     /**

core/src/main/java/ai/timefold/solver/core/api/score/stream/tri/TriConstraintStream.java

@@ -8,6 +8,7 @@
 import static ai.timefold.solver.core.impl.util.ConstantLambdaUtils.uniConstantNull;
 
 import java.math.BigDecimal;
+import java.util.Collection;
 import java.util.Objects;
 import java.util.function.BiFunction;
 import java.util.function.Function;
@@ -1094,6 +1095,19 @@ public interface TriConstraintStream<A, B, C> extends ConstraintStream {
             @NonNull TriFunction<A, B, C, ResultA_> mappingA, @NonNull TriFunction<A, B, C, ResultB_> mappingB,
             @NonNull TriFunction<A, B, C, ResultC_> mappingC, @NonNull TriFunction<A, B, C, ResultD_> mappingD);
 
+    /**
+     * As defined by {@link BiConstraintStream#flatten(BiFunction)}.
+     *
+     * @param <ResultD_> the type of the last fact in the resulting tuples.
+     *        It is recommended that this type be deeply immutable.
+     *        Not following this recommendation may lead to hard-to-debug hashing issues down the stream,
+     *        especially if this value is ever used as a group key.
+     * @param mapping function to convert the original tuple into {@link Iterable}.
+     *        For performance, returning an implementation of {@link Collection} is preferred.
+     */
+    <ResultD_> @NonNull QuadConstraintStream<A, B, C, ResultD_>
+            flatten(@NonNull TriFunction<A, B, C, Iterable<ResultD_>> mapping);
+
     /**
      * As defined by {@link BiConstraintStream#flattenLast(Function)}.
      *
@@ -1102,7 +1116,7 @@ public interface TriConstraintStream<A, B, C> extends ConstraintStream {
      *        Not following this recommendation may lead to hard-to-debug hashing issues down the stream,
      *        especially if this value is ever used as a group key.
      * @param mapping function to convert the last fact in the original tuple into {@link Iterable}.
-     *        For performance, returning an implementation of {@link java.util.Collection} is preferred.
+     *        For performance, returning an implementation of {@link Collection} is preferred.
      */
     <ResultC_> @NonNull TriConstraintStream<A, B, ResultC_> flattenLast(@NonNull Function<C, Iterable<ResultC_>> mapping);
 

core/src/main/java/ai/timefold/solver/core/api/score/stream/uni/UniConstraintStream.java

@@ -8,6 +8,7 @@
 
 import java.math.BigDecimal;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Objects;
 import java.util.function.BiPredicate;
 import java.util.function.Function;
@@ -1497,27 +1498,37 @@ default UniConstraintStream<A> ifNotExistsOtherIncludingNullVars(Class<A> otherC
 
     /**
      * Takes each tuple and applies a mapping on it, which turns the tuple into a {@link Iterable}.
-     * Returns a constraint stream consisting of contents of those iterables.
+     * Returns a constraint stream consisting of new tuples,
+     * each made of the original fact and one item from that iterable.
      * This may produce a stream with duplicate tuples.
      * See {@link #distinct()} for details.
      *
      * <p>
-     * In cases where the original tuple is already an {@link Iterable},
-     * use {@link Function#identity()} as the argument.
-     *
-     * <p>
      * Simple example: assuming a constraint stream of tuples of {@code Person}s
      * {@code [Ann(roles = [USER, ADMIN]]), Beth(roles = [USER]), Cathy(roles = [ADMIN, AUDITOR])]},
-     * calling {@code flattenLast(Person::getRoles)} on such stream will produce
-     * a stream of {@code [USER, ADMIN, USER, ADMIN, AUDITOR]}.
+     * calling {@code flatten(Person::getRoles)} on such stream will produce
+     * a stream of {@code [(Ann, USER), (Ann, ADMIN), (Beth, USER), (Cathy, ADMIN), (Cathy, AUDITOR)]}.
      *
      * @param mapping function to convert the original tuple into {@link Iterable}.
-     *        For performance, returning an implementation of {@link java.util.Collection} is preferred.
-     * @param <ResultA_> the type of facts in the resulting tuples.
+     *        For performance, returning an implementation of {@link Collection} is preferred.
+     * @param <ResultA_> the type of the last fact in the resulting tuples.
      *        It is recommended that this type be deeply immutable.
      *        Not following this recommendation may lead to hard-to-debug hashing issues down the stream,
      *        especially if this value is ever used as a group key.
      */
+    <ResultA_> @NonNull BiConstraintStream<A, ResultA_> flatten(@NonNull Function<A, Iterable<ResultA_>> mapping);
+
+    /**
+     * As defined by {@link #flatten(Function)},
+     * only replacing the only fact in the original tuple by an item from the iterable.
+     * This means the resulting stream is a {@link UniConstraintStream},
+     * not a {@link BiConstraintStream}.
+     * <p>
+     * Simple example: assuming a constraint stream of tuples of {@code Person}s
+     * {@code [Ann(roles = [USER, ADMIN]]), Beth(roles = [USER]), Cathy(roles = [ADMIN, AUDITOR])]},
+     * calling {@code flattenLast(Person::getRoles)} on such stream will produce
+     * a stream of {@code [(USER), (ADMIN), (USER), (ADMIN), (AUDITOR)]}.
+     */
     <ResultA_> @NonNull UniConstraintStream<ResultA_> flattenLast(@NonNull Function<A, Iterable<ResultA_>> mapping);
 
     /**

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

@@ -0,0 +1,33 @@
+package ai.timefold.solver.core.impl.bavet.bi;
+
+import java.util.Objects;
+import java.util.function.BiFunction;
+
+import ai.timefold.solver.core.impl.bavet.common.AbstractFlattenNode;
+import ai.timefold.solver.core.impl.bavet.common.tuple.BiTuple;
+import ai.timefold.solver.core.impl.bavet.common.tuple.TriTuple;
+import ai.timefold.solver.core.impl.bavet.common.tuple.TupleLifecycle;
+
+public final class FlattenBiNode<A, B, NewC> extends AbstractFlattenNode<BiTuple<A, B>, TriTuple<A, B, NewC>, NewC> {
+
+    private final BiFunction<A, B, Iterable<NewC>> mappingFunction;
+    private final int outputStoreSize;
+
+    public FlattenBiNode(int flattenStoreIndex, BiFunction<A, B, Iterable<NewC>> mappingFunction,
+            TupleLifecycle<TriTuple<A, B, NewC>> nextNodesTupleLifecycle, int outputStoreSize) {
+        super(flattenStoreIndex, nextNodesTupleLifecycle);
+        this.mappingFunction = Objects.requireNonNull(mappingFunction);
+        this.outputStoreSize = outputStoreSize;
+    }
+
+    @Override
+    protected TriTuple<A, B, NewC> createTuple(BiTuple<A, B> originalTuple, NewC newC) {
+        return TriTuple.of(originalTuple.getA(), originalTuple.getB(), newC, outputStoreSize);
+    }
+
+    @Override
+    protected Iterable<NewC> extractIterable(BiTuple<A, B> tuple) {
+        return mappingFunction.apply(tuple.getA(), tuple.getB());
+    }
+
+}

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

@@ -1,18 +1,21 @@
 package ai.timefold.solver.core.impl.bavet.bi;
 
+import java.util.Objects;
 import java.util.function.Function;
 
-import ai.timefold.solver.core.impl.bavet.common.AbstractFlattenLastNode;
+import ai.timefold.solver.core.impl.bavet.common.AbstractFlattenNode;
 import ai.timefold.solver.core.impl.bavet.common.tuple.BiTuple;
 import ai.timefold.solver.core.impl.bavet.common.tuple.TupleLifecycle;
 
-public final class FlattenLastBiNode<A, B, NewB> extends AbstractFlattenLastNode<BiTuple<A, B>, BiTuple<A, NewB>, B, NewB> {
+public final class FlattenLastBiNode<A, B, NewB> extends AbstractFlattenNode<BiTuple<A, B>, BiTuple<A, NewB>, NewB> {
 
+    private final Function<B, Iterable<NewB>> mappingFunction;
     private final int outputStoreSize;
 
     public FlattenLastBiNode(int flattenLastStoreIndex, Function<B, Iterable<NewB>> mappingFunction,
             TupleLifecycle<BiTuple<A, NewB>> nextNodesTupleLifecycle, int outputStoreSize) {
-        super(flattenLastStoreIndex, mappingFunction, nextNodesTupleLifecycle);
+        super(flattenLastStoreIndex, nextNodesTupleLifecycle);
+        this.mappingFunction = Objects.requireNonNull(mappingFunction);
         this.outputStoreSize = outputStoreSize;
     }
 
@@ -22,8 +25,8 @@ protected BiTuple<A, NewB> createTuple(BiTuple<A, B> originalTuple, NewB newB) {
     }
 
     @Override
-    protected B getEffectiveFactIn(BiTuple<A, B> tuple) {
-        return tuple.getB();
+    protected Iterable<NewB> extractIterable(BiTuple<A, B> tuple) {
+        return mappingFunction.apply(tuple.getB());
     }
 
 }

core/src/main/java/ai/timefold/solver/core/impl/bavet/common/AbstractFlattenLastNode.java → core/src/main/java/ai/timefold/solver/core/impl/bavet/common/AbstractFlattenNode.java

@@ -7,38 +7,33 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.function.Consumer;
-import java.util.function.Function;
 import java.util.function.Supplier;
 
 import ai.timefold.solver.core.impl.bavet.common.tuple.Tuple;
 import ai.timefold.solver.core.impl.bavet.common.tuple.TupleLifecycle;
 import ai.timefold.solver.core.impl.bavet.common.tuple.TupleState;
 import ai.timefold.solver.core.impl.util.CollectionUtils;
 
-public abstract class AbstractFlattenLastNode<InTuple_ extends Tuple, OutTuple_ extends Tuple, EffectiveItem_, FlattenedItem_>
+public abstract class AbstractFlattenNode<InTuple_ extends Tuple, OutTuple_ extends Tuple, FlattenedItem_>
         extends AbstractNode
         implements TupleLifecycle<InTuple_> {
 
-    private final int flattenLastStoreIndex;
-    private final Function<EffectiveItem_, Iterable<FlattenedItem_>> mappingFunction;
+    private final int flattenStoreIndex;
     private final StaticPropagationQueue<OutTuple_> propagationQueue;
 
-    protected AbstractFlattenLastNode(int flattenLastStoreIndex,
-            Function<EffectiveItem_, Iterable<FlattenedItem_>> mappingFunction,
-            TupleLifecycle<OutTuple_> nextNodesTupleLifecycle) {
-        this.flattenLastStoreIndex = flattenLastStoreIndex;
-        this.mappingFunction = Objects.requireNonNull(mappingFunction);
+    protected AbstractFlattenNode(int flattenStoreIndex, TupleLifecycle<OutTuple_> nextNodesTupleLifecycle) {
+        this.flattenStoreIndex = flattenStoreIndex;
         this.propagationQueue = new StaticPropagationQueue<>(nextNodesTupleLifecycle);
     }
 
     @Override
     public final void insert(InTuple_ tuple) {
-        if (tuple.getStore(flattenLastStoreIndex) != null) {
+        if (tuple.getStore(flattenStoreIndex) != null) {
             throw new IllegalStateException(
                     "Impossible state: the input for the tuple (%s) was already added in the tupleStore."
                             .formatted(tuple));
         }
-        var iterable = mappingFunction.apply(getEffectiveFactIn(tuple));
+        var iterable = extractIterable(tuple);
         if (iterable instanceof Collection<FlattenedItem_> collection) {
             // Optimization for Collection, where we know the size.
             var size = collection.size();
@@ -49,7 +44,7 @@ public final void insert(InTuple_ tuple) {
             for (var item : collection) {
                 addTuple(tuple, item, bagByItem);
             }
-            tuple.setStore(flattenLastStoreIndex, bagByItem);
+            tuple.setStore(flattenStoreIndex, bagByItem);
         } else {
             var iterator = iterable.iterator();
             if (!iterator.hasNext()) {
@@ -59,7 +54,7 @@ public final void insert(InTuple_ tuple) {
             while (iterator.hasNext()) {
                 addTuple(tuple, iterator.next(), bagByItem);
             }
-            tuple.setStore(flattenLastStoreIndex, bagByItem);
+            tuple.setStore(flattenStoreIndex, bagByItem);
         }
     }
 
@@ -75,26 +70,26 @@ private void addTuple(InTuple_ originalTuple, FlattenedItem_ item,
 
     @Override
     public final void update(InTuple_ tuple) {
-        FlattenBagByItem<FlattenedItem_, OutTuple_> bagByItem = tuple.getStore(flattenLastStoreIndex);
+        FlattenBagByItem<FlattenedItem_, OutTuple_> bagByItem = tuple.getStore(flattenStoreIndex);
         if (bagByItem == null) {
             // No fail fast if null because we don't track which tuples made it through the filter predicate(s).
             insert(tuple);
             return;
         }
 
         bagByItem.resetAll();
-        for (var item : mappingFunction.apply(getEffectiveFactIn(tuple))) {
+        for (var item : extractIterable(tuple)) {
             addTuple(tuple, item, bagByItem);
         }
         bagByItem.getAllBags()
                 .removeIf(bag -> bag.removeExtras(this::removeTuple));
     }
 
-    protected abstract EffectiveItem_ getEffectiveFactIn(InTuple_ tuple);
+    protected abstract Iterable<FlattenedItem_> extractIterable(InTuple_ tuple);
 
     @Override
     public final void retract(InTuple_ tuple) {
-        FlattenBagByItem<FlattenedItem_, OutTuple_> bagByItem = tuple.removeStore(flattenLastStoreIndex);
+        FlattenBagByItem<FlattenedItem_, OutTuple_> bagByItem = tuple.removeStore(flattenStoreIndex);
         if (bagByItem == null) {
             // No fail fast if null because we don't track which tuples made it through the filter predicate(s)
             return;

core/src/main/java/ai/timefold/solver/core/impl/bavet/quad/FlattenLastQuadNode.java

@@ -1,19 +1,22 @@
 package ai.timefold.solver.core.impl.bavet.quad;
 
+import java.util.Objects;
 import java.util.function.Function;
 
-import ai.timefold.solver.core.impl.bavet.common.AbstractFlattenLastNode;
+import ai.timefold.solver.core.impl.bavet.common.AbstractFlattenNode;
 import ai.timefold.solver.core.impl.bavet.common.tuple.QuadTuple;
 import ai.timefold.solver.core.impl.bavet.common.tuple.TupleLifecycle;
 
 public final class FlattenLastQuadNode<A, B, C, D, NewD>
-        extends AbstractFlattenLastNode<QuadTuple<A, B, C, D>, QuadTuple<A, B, C, NewD>, D, NewD> {
+        extends AbstractFlattenNode<QuadTuple<A, B, C, D>, QuadTuple<A, B, C, NewD>, NewD> {
 
+    private final Function<D, Iterable<NewD>> mappingFunction;
     private final int outputStoreSize;
 
     public FlattenLastQuadNode(int flattenLastStoreIndex, Function<D, Iterable<NewD>> mappingFunction,
             TupleLifecycle<QuadTuple<A, B, C, NewD>> nextNodesTupleLifecycle, int outputStoreSize) {
-        super(flattenLastStoreIndex, mappingFunction, nextNodesTupleLifecycle);
+        super(flattenLastStoreIndex, nextNodesTupleLifecycle);
+        this.mappingFunction = Objects.requireNonNull(mappingFunction);
         this.outputStoreSize = outputStoreSize;
     }
 
@@ -23,8 +26,8 @@ protected QuadTuple<A, B, C, NewD> createTuple(QuadTuple<A, B, C, D> originalTup
     }
 
     @Override
-    protected D getEffectiveFactIn(QuadTuple<A, B, C, D> tuple) {
-        return tuple.getD();
+    protected Iterable<NewD> extractIterable(QuadTuple<A, B, C, D> tuple) {
+        return mappingFunction.apply(tuple.getD());
     }
 
 }