fix: iterators must not survive step boundaries

Author: triceo

core/src/main/java/ai/timefold/solver/core/impl/neighborhood/stream/BiRandomMoveIterator.java

@@ -1,7 +1,9 @@
 package ai.timefold.solver.core.impl.neighborhood.stream;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.Iterator;
+import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.random.RandomGenerator;
@@ -60,18 +62,23 @@ final class BiRandomMoveIterator<Solution_, A, B> implements Iterator<Move<Solut
 
     private final BiMoveStreamContext<Solution_, A, B> context;
     private final RandomGenerator workingRandom;
+    /**
+     * We cannot cache the right iterator on the left tuple,
+     * because that cache would survive across steps and the right iterator would be invalid by then.
+     * The alternative would be worse than this map,
+     * likely requiring some loop over tuples at step end, clearing the cache.
+     */
+    private final Map<UniTuple<A>, Iterator<UniTuple<B>>> leftTupleToRightIteratorMap = new HashMap<>();
 
     // Fields required for iteration.
     private final Iterator<UniTuple<A>> leftTupleIterator;
-    private final int rightIteratorStoreIndex;
     private @Nullable Move<Solution_> nextMove;
 
     public BiRandomMoveIterator(BiMoveStreamContext<Solution_, A, B> context, RandomGenerator workingRandom) {
         this.context = Objects.requireNonNull(context);
         this.workingRandom = Objects.requireNonNull(workingRandom);
-        var leftDatasetInstance = context.getLeftDatasetInstance();
-        this.rightIteratorStoreIndex = leftDatasetInstance.getRightIteratorStoreIndex();
-        this.leftTupleIterator = leftDatasetInstance.randomIterator(workingRandom);
+        this.leftTupleIterator = context.getLeftDatasetInstance()
+                .randomIterator(workingRandom);
     }
 
     @Override
@@ -82,10 +89,9 @@ public boolean hasNext() {
 
         while (leftTupleIterator.hasNext()) {
             var leftTuple = leftTupleIterator.next();
-            var rightEmpty = pickNextMove(leftTuple);
-            if (rightEmpty) {
+            if (!pickNextMove(leftTuple)) {
                 leftTupleIterator.remove();
-                leftTuple.setStore(rightIteratorStoreIndex, null);
+                leftTupleToRightIteratorMap.remove(leftTuple);
             }
             if (nextMove != null) {
                 if (nextMove instanceof AbstractSelectorBasedMove<Solution_> legacyMove) {
@@ -101,17 +107,20 @@ Please refactor your code (%s) to use the new Move API."""
     }
 
     private boolean pickNextMove(UniTuple<A> leftTuple) {
-        var rightTupleIterator = (Iterator<UniTuple<B>>) leftTuple.getStore(rightIteratorStoreIndex);
+        var rightTupleIterator = leftTupleToRightIteratorMap.get(leftTuple);
         if (rightTupleIterator == null) {
             rightTupleIterator = createRightTupleIterator(leftTuple);
-            leftTuple.setStore(rightIteratorStoreIndex, rightTupleIterator);
+            if (!rightTupleIterator.hasNext()) {
+                return false;
+            }
+            leftTupleToRightIteratorMap.put(leftTuple, rightTupleIterator);
         }
         if (!rightTupleIterator.hasNext()) {
-            return true;
+            return false;
         }
         nextMove = context.buildMove(leftTuple.getA(), rightTupleIterator.next().getA());
         rightTupleIterator.remove();
-        return false;
+        return true;
     }
 
     private Iterator<UniTuple<B>> createRightTupleIterator(UniTuple<A> leftTuple) {

core/src/main/java/ai/timefold/solver/core/impl/neighborhood/stream/enumerating/common/AbstractLeftDataset.java

@@ -12,7 +12,6 @@ protected AbstractLeftDataset(AbstractUniEnumeratingStream<Solution_, A> parent)
         super(parent);
     }
 
-    public abstract AbstractLeftDatasetInstance<Solution_, UniTuple<A>> instantiate(int rightSequenceStoreIndex,
-            int entryStoreIndex);
+    public abstract AbstractLeftDatasetInstance<Solution_, UniTuple<A>> instantiate(int entryStoreIndex);
 
 }

core/src/main/java/ai/timefold/solver/core/impl/neighborhood/stream/enumerating/common/AbstractLeftDatasetInstance.java

@@ -15,15 +15,9 @@ public abstract class AbstractLeftDatasetInstance<Solution_, Tuple_ extends Tupl
         implements Iterable<Tuple_> {
 
     private final ElementAwareArrayList<Tuple_> tupleList = new ElementAwareArrayList<>();
-    private final int rightIteratorStoreIndex;
 
-    protected AbstractLeftDatasetInstance(AbstractDataset<Solution_> parent, int rightIteratorStoreIndex, int entryStoreIndex) {
+    protected AbstractLeftDatasetInstance(AbstractDataset<Solution_> parent, int entryStoreIndex) {
         super(parent, entryStoreIndex);
-        this.rightIteratorStoreIndex = rightIteratorStoreIndex;
-    }
-
-    public int getRightIteratorStoreIndex() {
-        return rightIteratorStoreIndex;
     }
 
     @Override

core/src/main/java/ai/timefold/solver/core/impl/neighborhood/stream/enumerating/uni/LeftTerminalUniEnumeratingStream.java

@@ -22,9 +22,7 @@ public LeftTerminalUniEnumeratingStream(EnumeratingStreamFactory<Solution_> enum
     @Override
     public void buildNode(DataNodeBuildHelper<Solution_> buildHelper) {
         assertEmptyChildStreamList();
-        var datasetInstance = dataset.instantiate(
-                buildHelper.reserveTupleStoreIndex(parent.getTupleSource()),
-                buildHelper.reserveTupleStoreIndex(parent.getTupleSource()));
+        var datasetInstance = dataset.instantiate(buildHelper.reserveTupleStoreIndex(parent.getTupleSource()));
         buildHelper.putInsertUpdateRetract(this, datasetInstance);
     }
 

core/src/main/java/ai/timefold/solver/core/impl/neighborhood/stream/enumerating/uni/UniLeftDataset.java

@@ -14,8 +14,8 @@ public UniLeftDataset(AbstractUniEnumeratingStream<Solution_, A> parent) {
     }
 
     @Override
-    public UniLeftDatasetInstance<Solution_, A> instantiate(int rightIteratorStoreIndex, int entryStoreIndex) {
-        return new UniLeftDatasetInstance<>(this, rightIteratorStoreIndex, entryStoreIndex);
+    public UniLeftDatasetInstance<Solution_, A> instantiate(int entryStoreIndex) {
+        return new UniLeftDatasetInstance<>(this, entryStoreIndex);
     }
 
     @Override

core/src/main/java/ai/timefold/solver/core/impl/neighborhood/stream/enumerating/uni/UniLeftDatasetInstance.java

@@ -10,8 +10,8 @@
 public final class UniLeftDatasetInstance<Solution_, A>
         extends AbstractLeftDatasetInstance<Solution_, UniTuple<A>> {
 
-    public UniLeftDatasetInstance(AbstractDataset<Solution_> parent, int rightIteratorStoreIndex, int entryStoreIndex) {
-        super(parent, rightIteratorStoreIndex, entryStoreIndex);
+    public UniLeftDatasetInstance(AbstractDataset<Solution_> parent, int entryStoreIndex) {
+        super(parent, entryStoreIndex);
     }
 
 }