Optimizations around the counters

Author: triceo

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

@@ -42,8 +42,7 @@ protected AbstractIndexedIfExistsNode(boolean shouldExist,
         this.inputStoreIndexLeftKeys = leftTupleStorePositionTracker.reserveNextAvailablePosition();
         this.inputStoreIndexLeftCounter = leftTupleStorePositionTracker.reserveNextAvailablePosition();
         this.inputStoreIndexRightKeys = rightTupleStorePositionTracker.reserveNextAvailablePosition();
-        this.indexerLeft = indexerFactory.buildIndexer(true,
-                new ExistsCounterPositionTracker<>(leftTupleStorePositionTracker.reserveNextAvailablePosition()));
+        this.indexerLeft = indexerFactory.buildIndexer(true, ExistsCounterPositionTracker.instance());
         this.indexerRight = indexerFactory.buildIndexer(false,
                 new TuplePositionTracker<>(rightTupleStorePositionTracker.reserveNextAvailablePosition()));
     }

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

@@ -32,7 +32,7 @@ protected AbstractUnindexedIfExistsNode(boolean shouldExist, TupleStorePositionT
         this.inputStoreIndexLeftCounter = leftTupleStorePositionTracker.reserveNextAvailablePosition();
         this.inputStoreIndexRightTuple = rightTupleStorePositionTracker.reserveNextAvailablePosition();
         this.leftCounterSet = new IndexedSet<>(
-                new ExistsCounterPositionTracker<>(leftTupleStorePositionTracker.reserveNextAvailablePosition()));
+                new ExistsCounterPositionTracker<>());
         this.rightTupleSet = new IndexedSet<>(new TuplePositionTracker<>(inputStoreIndexRightTuple));
     }
 

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

@@ -3,12 +3,16 @@
 import ai.timefold.solver.core.impl.bavet.common.tuple.AbstractTuple;
 import ai.timefold.solver.core.impl.bavet.common.tuple.TupleState;
 
+import org.jspecify.annotations.NullMarked;
+
+@NullMarked
 public final class ExistsCounter<Tuple_ extends AbstractTuple>
         extends AbstractPropagationMetadataCarrier<Tuple_> {
 
     final Tuple_ leftTuple;
     TupleState state = TupleState.DEAD; // It's the node's job to mark a new instance as CREATING.
     int countRight = 0;
+    int indexedSetPositon = -1;
 
     ExistsCounter(Tuple_ leftTuple) {
         this.leftTuple = leftTuple;

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

@@ -3,6 +3,8 @@
 import ai.timefold.solver.core.impl.bavet.common.index.IndexedSet;
 import ai.timefold.solver.core.impl.bavet.common.tuple.AbstractTuple;
 
+import org.jspecify.annotations.NullMarked;
+
 /**
  * Used for filtering in {@link AbstractIfExistsNode}.
  * There is no place where both left and right sets for each counter would be kept together,
@@ -11,6 +13,7 @@
  *
  * @param <LeftTuple_>
  */
+@NullMarked
 final class ExistsCounterHandle<LeftTuple_ extends AbstractTuple> {
 
     final ExistsCounter<LeftTuple_> counter;

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

@@ -1,37 +1,29 @@
 package ai.timefold.solver.core.impl.bavet.common;
 
+import java.util.function.ToIntFunction;
+
 import ai.timefold.solver.core.impl.bavet.common.index.ElementPositionTracker;
 import ai.timefold.solver.core.impl.bavet.common.tuple.AbstractTuple;
 
 import org.jspecify.annotations.NullMarked;
 
 @SuppressWarnings({ "rawtypes", "unchecked" })
 @NullMarked
-record ExistsCounterHandlePositionTracker<Tuple_ extends AbstractTuple>(PositionGetter<Tuple_> positionGetter,
-        PositionClearer<Tuple_> positionClearer,
+record ExistsCounterHandlePositionTracker<Tuple_ extends AbstractTuple>(
+        ToIntFunction<ExistsCounterHandle<Tuple_>> positionGetter,
         PositionSetter<Tuple_> positionSetter)
         implements
             ElementPositionTracker<ExistsCounterHandle<Tuple_>> {
 
     private static final ExistsCounterHandlePositionTracker LEFT = new ExistsCounterHandlePositionTracker(
-            tracker -> tracker.leftPosition,
-            tracker -> {
-                var result = tracker.leftPosition;
-                tracker.leftPosition = -1;
-                return result;
-            },
+            (ToIntFunction<ExistsCounterHandle>) tracker -> tracker.leftPosition,
             (tracker, position) -> {
                 var oldValue = tracker.leftPosition;
                 tracker.leftPosition = position;
                 return oldValue;
             });
     private static final ExistsCounterHandlePositionTracker RIGHT = new ExistsCounterHandlePositionTracker(
-            tracker -> tracker.rightPosition,
-            tracker -> {
-                var result = tracker.rightPosition;
-                tracker.rightPosition = -1;
-                return result;
-            },
+            (ToIntFunction<ExistsCounterHandle>) tracker -> tracker.rightPosition,
             (tracker, position) -> {
                 var oldValue = tracker.rightPosition;
                 tracker.rightPosition = position;
@@ -47,34 +39,15 @@ public static <Tuple_ extends AbstractTuple> ExistsCounterHandlePositionTracker<
     }
 
     @Override
-    public int getPosition(ExistsCounterHandle<Tuple_> element) {
-        return positionGetter.apply(element);
-    }
-
-    @Override
-    public int setPosition(ExistsCounterHandle<Tuple_> element, int position) {
-        return positionSetter.apply(element, position);
+    public void setPosition(ExistsCounterHandle<Tuple_> element, int position) {
+        positionSetter.apply(element, position);
     }
 
     @Override
     public int clearPosition(ExistsCounterHandle<Tuple_> element) {
-        return positionClearer.apply(element);
-    }
-
-    @FunctionalInterface
-    @NullMarked
-    interface PositionGetter<Tuple_ extends AbstractTuple> {
-
-        int apply(ExistsCounterHandle<Tuple_> tracker);
-
-    }
-
-    @FunctionalInterface
-    @NullMarked
-    interface PositionClearer<Tuple_ extends AbstractTuple> {
-
-        int apply(ExistsCounterHandle<Tuple_> tracker);
-
+        var oldPosition = positionGetter.applyAsInt(element);
+        positionSetter.apply(element, -1);
+        return oldPosition;
     }
 
     @FunctionalInterface

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

@@ -3,32 +3,30 @@
 import ai.timefold.solver.core.impl.bavet.common.index.ElementPositionTracker;
 import ai.timefold.solver.core.impl.bavet.common.tuple.AbstractTuple;
 
-record ExistsCounterPositionTracker<Tuple_ extends AbstractTuple>(int inputStorePosition)
+import org.jspecify.annotations.NullMarked;
+
+@SuppressWarnings({ "unchecked", "rawtypes" })
+@NullMarked
+record ExistsCounterPositionTracker<Tuple_ extends AbstractTuple>()
         implements
             ElementPositionTracker<ExistsCounter<Tuple_>> {
 
-    @Override
-    public int getPosition(ExistsCounter<Tuple_> element) {
-        var tuple = element.getTuple();
-        var value = tuple.getStore(inputStorePosition);
-        return value == null ? -1 : (int) value;
+    private static final ExistsCounterPositionTracker INSTANCE = new ExistsCounterPositionTracker();
+
+    public static <Tuple_ extends AbstractTuple> ExistsCounterPositionTracker<Tuple_> instance() {
+        return INSTANCE;
     }
 
     @Override
-    public int setPosition(ExistsCounter<Tuple_> element, int position) {
-        var tuple = element.getTuple();
-        var oldValue = getPosition(element);
-        tuple.setStore(inputStorePosition, position);
-        return oldValue;
+    public void setPosition(ExistsCounter<Tuple_> element, int position) {
+        element.indexedSetPositon = position;
     }
 
     @Override
     public int clearPosition(ExistsCounter<Tuple_> element) {
-        try {
-            return element.getTuple().removeStore(inputStorePosition);
-        } catch (NullPointerException e) {
-            return -1;
-        }
+        var oldPosition = element.indexedSetPositon;
+        element.indexedSetPositon = -1;
+        return oldPosition;
     }
 
 }

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

@@ -3,21 +3,16 @@
 import ai.timefold.solver.core.impl.bavet.common.index.ElementPositionTracker;
 import ai.timefold.solver.core.impl.bavet.common.tuple.AbstractTuple;
 
+import org.jspecify.annotations.NullMarked;
+
+@NullMarked
 public record TuplePositionTracker<Tuple_ extends AbstractTuple>(int inputStorePosition)
         implements
             ElementPositionTracker<Tuple_> {
 
     @Override
-    public int getPosition(Tuple_ element) {
-        var value = element.getStore(inputStorePosition);
-        return value == null ? -1 : (int) value;
-    }
-
-    @Override
-    public int setPosition(Tuple_ element, int position) {
-        var oldValue = getPosition(element);
+    public void setPosition(Tuple_ element, int position) {
         element.setStore(inputStorePosition, position);
-        return oldValue;
     }
 
     @Override

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

@@ -1,29 +1,23 @@
 package ai.timefold.solver.core.impl.bavet.common.index;
 
+import org.jspecify.annotations.NullMarked;
+
 /**
  * Allows to read and modify the position of an element in an {@link IndexedSet}.
  * Typically points to a field in the element itself.
  *
  * @param <T>
  */
+@NullMarked
 public interface ElementPositionTracker<T> {
 
-    /**
-     * Gets the position of the given element.
-     *
-     * @param element never null
-     * @return >= 0 if the element is tracked, or -1 if it is not tracked
-     */
-    int getPosition(T element);
-
     /**
      * Sets the position of the given element.
      *
      * @param element never null
      * @param position >= 0
-     * @return the previous position of the element, or -1 if it was not tracked before
      */
-    int setPosition(T element, int position);
+    void setPosition(T element, int position);
 
     /**
      * Clears the position of the given element.

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

@@ -12,8 +12,9 @@
 import org.jspecify.annotations.Nullable;
 
 /**
- * {@link ArrayList}-backed set which allows to {@link #remove(Object)} an element without knowing its position
- * and without an expensive lookup.
+ * {@link ArrayList}-backed set which allows to {@link #remove(Object)} an element
+ * without knowing its position and without an expensive lookup.
+ * It also allows for direct random access like a list.
  * <p>
  * It uses an {@link ElementPositionTracker} to track the insertion position of each element.
  * When an element is removed, the insertion position of later elements is not changed.
@@ -29,8 +30,10 @@
  * Random access is not required for Constraint Streams, but Neighborhoods make heavy use of it;
  * if we used the {@link ElementAwareList} implementation instead,
  * we would have to copy the elements to an array every time we need to access them randomly during move generation.
- * 
- * 
+ * <p>
+ * For performance reasons, this class does not check if an element was already added;
+ * duplicates must be avoided by the caller and will cause undefined behavior.
+ *
  * @param <T>
  */
 @NullMarked
@@ -52,20 +55,19 @@ private List<T> getElementList() {
     }
 
     /**
-     * Appends the specified element to the end of this collection, if not already present.
-     * Will use identity comparison to check for presence;
-     * two different instances which {@link Object#equals(Object) equal} are considered different elements.
+     * Appends the specified element to the end of this collection.
+     * If the element is already present,
+     * undefined, unexpected, and incorrect behavior should be expected.
+     * <p>
+     * Presence of the element can be checked using the associated {@link ElementPositionTracker}.
+     * For performance reasons, this method avoids that check.
      *
      * @param element element to be appended to this collection
-     * @throws IllegalStateException if the element was already present in this collection
      */
     public void add(T element) {
         var actualElementList = getElementList();
         actualElementList.add(element);
-        if (elementPositionTracker.setPosition(element, actualElementList.size() - 1) >= 0) {
-            throw new IllegalStateException("Impossible state: the element (%s) was already added to the IndexedSet."
-                    .formatted(element));
-        }
+        elementPositionTracker.setPosition(element, actualElementList.size() - 1);
     }
 
     /**

core/src/test/java/ai/timefold/solver/core/impl/bavet/common/index/IndexedSetTest.java

@@ -25,16 +25,6 @@ void addMultipleElements() {
         assertThat(set.asList()).containsExactly("A", "B", "C");
     }
 
-    @Test
-    void addDuplicateElement() {
-        var set = new IndexedSet<>(stringTracker);
-
-        set.add("A");
-        assertThatThrownBy(() -> set.add("A"))
-                .isInstanceOf(IllegalStateException.class)
-                .hasMessageContaining("was already added");
-    }
-
     @Test
     void removeElement() {
         var set = new IndexedSet<>(stringTracker);
@@ -226,13 +216,8 @@ private static final class SimpleTracker<T> implements ElementPositionTracker<T>
         private final Map<T, Integer> positions = new HashMap<>();
 
         @Override
-        public int setPosition(T element, int position) {
-            return positions.put(element, position) == null ? -1 : position;
-        }
-
-        @Override
-        public int getPosition(T element) {
-            return positions.getOrDefault(element, -1);
+        public void setPosition(T element, int position) {
+            positions.put(element, position);
         }
 
         @Override