Do the same for joins

Author: triceo

core/src/main/java/ai/timefold/solver/core/impl/bavet/common/AbstractIndexedJoinNode.java

@@ -84,7 +84,7 @@ public final void updateLeft(LeftTuple_ leftTuple) {
         } else {
             indexerLeft.remove(oldIndexKeys, leftTuple);
             IndexedSet<OutTuple_> outTupleSetLeft = leftTuple.getStore(inputStoreIndexLeftOutTupleSet);
-            outTupleSetLeft.forEach(this::retractOutTuple);
+            outTupleSetLeft.clear(this::retractOutTupleByLeft);
             // outTupleSetLeft is now empty, no need for leftTuple.setStore(...);
             indexAndPropagateLeft(leftTuple, newIndexKeys);
         }
@@ -105,7 +105,7 @@ public final void retractLeft(LeftTuple_ leftTuple) {
         }
         IndexedSet<OutTuple_> outTupleSetLeft = leftTuple.removeStore(inputStoreIndexLeftOutTupleSet);
         indexerLeft.remove(indexKeys, leftTuple);
-        outTupleSetLeft.forEach(this::retractOutTuple);
+        outTupleSetLeft.clear(this::retractOutTupleByLeft);
     }
 
     @Override
@@ -137,7 +137,7 @@ public final void updateRight(UniTuple<Right_> rightTuple) {
         } else {
             IndexedSet<OutTuple_> outTupleSetRight = rightTuple.getStore(inputStoreIndexRightOutTupleSet);
             indexerRight.remove(oldIndexKeys, rightTuple);
-            outTupleSetRight.forEach(this::retractOutTuple);
+            outTupleSetRight.clear(this::retractOutTupleByRight);
             // outTupleSetRight is now empty, no need for rightTuple.setStore(...);
             indexAndPropagateRight(rightTuple, newIndexKeys);
         }
@@ -158,7 +158,7 @@ public final void retractRight(UniTuple<Right_> rightTuple) {
         }
         IndexedSet<OutTuple_> outTupleSetRight = rightTuple.removeStore(inputStoreIndexRightOutTupleSet);
         indexerRight.remove(indexKeys, rightTuple);
-        outTupleSetRight.forEach(this::retractOutTuple);
+        outTupleSetRight.clear(this::retractOutTupleByRight);
     }
 
 }

core/src/main/java/ai/timefold/solver/core/impl/bavet/common/AbstractJoinNode.java

@@ -128,6 +128,22 @@ private void findAndProcessOutTuple(LeftTuple_ leftTuple, UniTuple<Right_> right
         }
     }
 
+    private void retractOutTuple(OutTuple_ outTuple) {
+        removeFromLeftOutSet(outTuple);
+        removeFromRightOutSet(outTuple);
+        propagateRetract(outTuple);
+    }
+
+    private void removeFromLeftOutSet(OutTuple_ outTuple) {
+        IndexedSet<OutTuple_> outSetLeft = outTuple.removeStore(outputStoreIndexLeftOutSet);
+        outSetLeft.remove(outTuple);
+    }
+
+    private void removeFromRightOutSet(OutTuple_ outTuple) {
+        IndexedSet<OutTuple_> outSetRight = outTuple.removeStore(outputStoreIndexRightOutSet);
+        outSetRight.remove(outTuple);
+    }
+
     private static <OutTuple_ extends AbstractTuple> OutTuple_ findOutTuple(IndexedSet<OutTuple_> sourceSet,
             IndexedSet<OutTuple_> referenceSet, int outputStoreIndex) {
         // Hack: outTuple has no left/right input tuple reference, use the left/right outSet reference instead.
@@ -184,11 +200,13 @@ protected final void innerUpdateRight(UniTuple<Right_> rightTuple, Consumer<Cons
         }
     }
 
-    protected final void retractOutTuple(OutTuple_ outTuple) {
-        IndexedSet<OutTuple_> outSetLeft = outTuple.removeStore(outputStoreIndexLeftOutSet);
-        outSetLeft.remove(outTuple);
-        IndexedSet<OutTuple_> outSetRight = outTuple.removeStore(outputStoreIndexRightOutSet);
-        outSetRight.remove(outTuple);
+    protected final void retractOutTupleByLeft(OutTuple_ outTuple) {
+        outTuple.removeStore(outputStoreIndexLeftOutSet); // The tuple will be removed from the set by the caller.
+        removeFromRightOutSet(outTuple);
+        propagateRetract(outTuple);
+    }
+
+    private void propagateRetract(OutTuple_ outTuple) {
         var state = outTuple.state;
         if (!state.isActive()) { // Impossible because they shouldn't linger in the indexes.
             throw new IllegalStateException("Impossible state: The tuple (%s) in node (%s) is in an unexpected state (%s)."
@@ -197,6 +215,12 @@ protected final void retractOutTuple(OutTuple_ outTuple) {
         propagationQueue.retract(outTuple, state == TupleState.CREATING ? TupleState.ABORTING : TupleState.DYING);
     }
 
+    protected final void retractOutTupleByRight(OutTuple_ outTuple) {
+        removeFromLeftOutSet(outTuple);
+        outTuple.removeStore(outputStoreIndexRightOutSet); // The tuple will be removed from the set by the caller.
+        propagateRetract(outTuple);
+    }
+
     @Override
     public Propagator getPropagator() {
         return propagationQueue;

core/src/main/java/ai/timefold/solver/core/impl/bavet/common/AbstractUnindexedJoinNode.java

@@ -71,7 +71,7 @@ public final void retractLeft(LeftTuple_ leftTuple) {
         }
         IndexedSet<OutTuple_> outTupleSetLeft = leftTuple.removeStore(inputStoreIndexLeftOutTupleSet);
         leftTupleSet.remove(leftTuple);
-        outTupleSetLeft.forEach(this::retractOutTuple);
+        outTupleSetLeft.clear(this::retractOutTupleByLeft);
     }
 
     @Override
@@ -105,7 +105,7 @@ public final void retractRight(UniTuple<Right_> rightTuple) {
         }
         IndexedSet<OutTuple_> outTupleSetRight = rightTuple.removeStore(inputStoreIndexRightOutTupleSet);
         rightTupleSet.remove(rightTuple);
-        outTupleSetRight.forEach(this::retractOutTuple);
+        outTupleSetRight.clear(this::retractOutTupleByRight);
     }
 
 }

core/src/main/java/ai/timefold/solver/core/impl/bavet/common/index/IndexedSet.java

@@ -275,6 +275,39 @@ private void forEachCompacting(Consumer<T> elementConsumer) {
         return null;
     }
 
+    /**
+     * Empties the collection.
+     * For each element being removed, the given consumer is called.
+     * This is a more efficient implementation than removing elements individually.
+     *
+     * @param elementConsumer the consumer to be called for each element being removed
+     */
+    public void clear(Consumer<T> elementConsumer) {
+        var nonGapCount = size();
+        if (nonGapCount == 0) {
+            forceClear();
+            return;
+        }
+        var oldLastElementPosition = lastElementPosition;
+        for (var i = 0; i <= oldLastElementPosition; i++) {
+            var element = elementList.get(i);
+            if (element == null) {
+                continue;
+            }
+            elementConsumer.accept(element);
+            if (lastElementPosition != oldLastElementPosition) {
+                throw new IllegalStateException("Impossible state: the IndexedSet was modified while being cleared.");
+            }
+            elementPositionTracker.clearPosition(element);
+            // We can stop early once all non-gap elements have been processed.
+            nonGapCount--;
+            if (nonGapCount == 0) {
+                break;
+            }
+        }
+        forceClear();
+    }
+
     /**
      * Returns a standard {@link List} view of this collection.
      * Users mustn't modify the returned list, as that'd also change the underlying data structure.
@@ -306,30 +339,4 @@ private void forceCompaction() {
         }
     }
 
-    public void clear(Consumer<T> elementConsumer) {
-        var nonGapCount = size();
-        if (nonGapCount == 0) {
-            forceClear();
-            return;
-        }
-        var oldLastElementPosition = lastElementPosition;
-        for (var i = 0; i <= oldLastElementPosition; i++) {
-            var element = elementList.get(i);
-            if (element == null) {
-                continue;
-            }
-            elementConsumer.accept(element);
-            if (lastElementPosition != oldLastElementPosition) {
-                throw new IllegalStateException("Impossible state: the IndexedSet was modified while being cleared.");
-            }
-            elementPositionTracker.clearPosition(element);
-            // We can stop early once all non-gap elements have been processed.
-            nonGapCount--;
-            if (nonGapCount == 0) {
-                break;
-            }
-        }
-        forceClear();
-    }
-
 }