More refactors

Author: triceo

core/src/main/java/ai/timefold/solver/core/impl/score/stream/collector/AbstractLongAverageSlot.java

@@ -24,16 +24,16 @@ public AbstractLongAverageSlot(State state) {
     protected void addMapped(long input) {
         cachedInput = input;
         state.count++;
-        state.sum += input;
+        state.sum = Math.addExact(state.sum, input);
     }
 
     protected void updateMapped(long input) {
-        state.sum += input - cachedInput;
+        state.sum += Math.subtractExact(input, cachedInput);
         cachedInput = input;
     }
 
     protected void removeMapped() {
         state.count--;
-        state.sum -= cachedInput;
+        state.sum = Math.subtractExact(state.sum, cachedInput);
     }
 }

core/src/main/java/ai/timefold/solver/core/impl/score/stream/collector/AbstractLongSumSlot.java

@@ -19,15 +19,15 @@ public AbstractLongSumSlot(State state) {
 
     protected void addMapped(long input) {
         cachedInput = input;
-        state.sum += input;
+        state.sum = Math.addExact(state.sum, input);
     }
 
     protected void updateMapped(long input) {
-        state.sum += input - cachedInput;
+        state.sum += Math.subtractExact(input, cachedInput);
         cachedInput = input;
     }
 
     protected void removeMapped() {
-        state.sum -= cachedInput;
+        state.sum = Math.subtractExact(state.sum, cachedInput);
     }
 }

core/src/main/java/ai/timefold/solver/core/impl/score/stream/collector/AbstractReferenceSumSlot.java

@@ -1,5 +1,6 @@
 package ai.timefold.solver.core.impl.score.stream.collector;
 
+import java.util.Objects;
 import java.util.function.BinaryOperator;
 
 import org.jspecify.annotations.Nullable;
@@ -35,7 +36,7 @@ protected void addMapped(Result_ input) {
     }
 
     protected void updateMapped(Result_ input) {
-        if (cachedValue == input) {
+        if (Objects.equals(cachedValue, input)) {
             return;
         }
         state.current = state.subtractor.apply(state.current, cachedValue);

core/src/main/java/ai/timefold/solver/core/impl/score/stream/collector/AbstractToMapSlot.java

@@ -1,7 +1,6 @@
 package ai.timefold.solver.core.impl.score.stream.collector;
 
 import java.util.Map;
-import java.util.Objects;
 import java.util.Set;
 import java.util.function.BinaryOperator;
 import java.util.function.IntFunction;
@@ -34,29 +33,26 @@ State<Key_, Value_, Set_, Result_> multiMapState(Supplier<Result_> resultSupplie
     }
 
     private final State<Key_, Value_, ResultValue_, Result_> state;
-    private @Nullable Key_ cachedKey;
-    private @Nullable Value_ cachedValue;
+    private @Nullable ToMapPerKeyCounter<Key_, Value_> cachedCounter;
+    private @Nullable CountHolder<Value_> cachedHolder;
 
     public AbstractToMapSlot(State<Key_, Value_, ResultValue_, Result_> state) {
         this.state = state;
     }
 
     protected void addMapped(Key_ key, Value_ value) {
-        this.cachedKey = key;
-        this.cachedValue = value;
         state.container.add(key, value);
+        cachedCounter = state.container.lastCounter();
+        cachedHolder = state.container.lastHolder();
     }
 
     protected void updateMapped(Key_ key, Value_ value) {
-        if (Objects.equals(cachedKey, key) && Objects.equals(cachedValue, value)) {
-            return;
-        }
-        state.container.update(cachedKey, cachedValue, key, value);
-        cachedKey = key;
-        cachedValue = value;
+        state.container.update(cachedCounter, cachedHolder, key, value);
+        cachedCounter = state.container.lastCounter();
+        cachedHolder = state.container.lastHolder();
     }
 
     protected void removeMapped() {
-        state.container.remove(cachedKey, cachedValue);
+        state.container.remove(cachedCounter, cachedHolder);
     }
 }

core/src/main/java/ai/timefold/solver/core/impl/score/stream/collector/CountHolder.java

@@ -0,0 +1,13 @@
+package ai.timefold.solver.core.impl.score.stream.collector;
+
+final class CountHolder<Value_> {
+
+    final Value_ value;
+    long count;
+
+    CountHolder(Value_ value) {
+        this.value = value;
+        this.count = 1;
+    }
+
+}

core/src/main/java/ai/timefold/solver/core/impl/score/stream/collector/ToMapPerKeyCounter.java

@@ -5,51 +5,46 @@
 import java.util.function.BinaryOperator;
 import java.util.stream.Stream;
 
-final class ToMapPerKeyCounter<Value_> {
-
-    private final Map<Value_, Long> counts = new LinkedHashMap<>();
-
-    public long add(Value_ value) {
-        return counts.compute(value, (k, currentCount) -> {
-            if (currentCount == null) {
-                return 1L;
-            } else {
-                return currentCount + 1;
-            }
-        });
+final class ToMapPerKeyCounter<Key_, Value_> {
+
+    final Key_ key;
+    private final Map<Value_, CountHolder<Value_>> counts = new LinkedHashMap<>(0);
+
+    ToMapPerKeyCounter(Key_ key) {
+        this.key = key;
     }
 
-    public long update(Value_ oldValue, Value_ newValue) {
-        var removedCount = remove(oldValue);
-        add(newValue);
-        return removedCount;
+    CountHolder<Value_> add(Value_ value) {
+        var holder = counts.get(value);
+        if (holder == null) {
+            holder = new CountHolder<>(value);
+            counts.put(value, holder);
+        } else {
+            holder.count++;
+        }
+        return holder;
     }
 
-    public long remove(Value_ value) {
-        var newCount = counts.compute(value, (k, currentCount) -> {
-            if (currentCount > 1L) {
-                return currentCount - 1;
-            } else {
-                return null;
-            }
-        });
-        return newCount == null ? 0L : newCount;
+    void decrement(CountHolder<Value_> holder) {
+        holder.count--;
+        if (holder.count == 0) {
+            counts.remove(holder.value);
+        }
     }
 
-    public Value_ merge(BinaryOperator<Value_> mergeFunction) {
-        // Rebuilding the value from the collection is not incremental.
-        // The impact is negligible, assuming there are not too many values for the same key.
-        return counts.entrySet()
+    Value_ merge(BinaryOperator<Value_> mergeFunction) {
+        return counts.values()
                 .stream()
-                .map(e -> Stream.generate(e::getKey)
-                        .limit(e.getValue())
+                .map(h -> Stream.generate(() -> h.value)
+                        .limit(h.count)
                         .reduce(mergeFunction)
-                        .orElseThrow(() -> new IllegalStateException("Impossible state: Should have had at least one value.")))
+                        .orElseThrow(
+                                () -> new IllegalStateException("Impossible state: Should have had at least one value.")))
                 .reduce(mergeFunction)
                 .orElseThrow(() -> new IllegalStateException("Impossible state: Should have had at least one value."));
     }
 
-    public boolean isEmpty() {
+    boolean isEmpty() {
         return counts.isEmpty();
     }
 

core/src/main/java/ai/timefold/solver/core/impl/score/stream/collector/ToMapResultContainer.java

@@ -2,14 +2,23 @@
 
 import java.util.Map;
 
+import org.jspecify.annotations.Nullable;
+
 sealed interface ToMapResultContainer<Key_, Value_, ResultValue_, Result_ extends Map<Key_, ResultValue_>>
         permits ToMultiMapResultContainer, ToSimpleMapResultContainer {
 
     void add(Key_ key, Value_ value);
 
-    void update(Key_ oldKey, Value_ oldValue, Key_ newKey, Value_ newValue);
+    void update(ToMapPerKeyCounter<Key_, Value_> counter, CountHolder<Value_> holder,
+            Key_ newKey, Value_ newValue);
+
+    void remove(ToMapPerKeyCounter<Key_, Value_> counter, CountHolder<Value_> holder);
+
+    @Nullable
+    ToMapPerKeyCounter<Key_, Value_> lastCounter();
 
-    void remove(Key_ key, Value_ value);
+    @Nullable
+    CountHolder<Value_> lastHolder();
 
     Result_ getResult();
 

core/src/main/java/ai/timefold/solver/core/impl/score/stream/collector/ToMultiMapResultContainer.java

@@ -7,55 +7,78 @@
 import java.util.function.IntFunction;
 import java.util.function.Supplier;
 
+import org.jspecify.annotations.Nullable;
+
 final class ToMultiMapResultContainer<Key_, Value_, Set_ extends Set<Value_>, Result_ extends Map<Key_, Set_>>
         implements ToMapResultContainer<Key_, Value_, Set_, Result_> {
 
     private final Supplier<Set_> setSupplier;
     private final Result_ result;
-    private final Map<Key_, ToMapPerKeyCounter<Value_>> valueCounts = new HashMap<>();
+    private final Map<Key_, ToMapPerKeyCounter<Key_, Value_>> valueCounts = new HashMap<>(0);
+    private @Nullable ToMapPerKeyCounter<Key_, Value_> lastCounter;
+    private @Nullable CountHolder<Value_> lastHolder;
 
-    public ToMultiMapResultContainer(Supplier<Result_> resultSupplier, IntFunction<Set_> setFunction) {
+    ToMultiMapResultContainer(Supplier<Result_> resultSupplier, IntFunction<Set_> setFunction) {
         var nonNullSetFunction = Objects.requireNonNull(setFunction);
         this.setSupplier = () -> nonNullSetFunction.apply(0);
         this.result = Objects.requireNonNull(resultSupplier).get();
     }
 
     @Override
     public void add(Key_ key, Value_ value) {
-        var counter = valueCounts.computeIfAbsent(key, k -> new ToMapPerKeyCounter<>());
-        counter.add(value);
-        result.computeIfAbsent(key, k -> setSupplier.get())
-                .add(value);
+        lastCounter = valueCounts.computeIfAbsent(key, ToMapPerKeyCounter::new);
+        lastHolder = lastCounter.add(value);
+        if (lastHolder.count == 1) {
+            result.computeIfAbsent(key, k -> setSupplier.get()).add(value);
+        }
     }
 
     @Override
-    public void update(Key_ oldKey, Value_ oldValue, Key_ newKey, Value_ newValue) {
-        if (Objects.equals(oldKey, newKey)) {
-            var counter = valueCounts.get(oldKey);
-            var removedCount = counter.update(oldValue, newValue);
-            if (removedCount == 0) {
-                result.get(oldKey).remove(oldValue);
+    public void update(ToMapPerKeyCounter<Key_, Value_> counter, CountHolder<Value_> holder,
+            Key_ newKey, Value_ newValue) {
+        if (Objects.equals(counter.key, newKey) && Objects.equals(holder.value, newValue)) {
+            lastCounter = counter;
+            lastHolder = holder;
+            return;
+        }
+        if (Objects.equals(counter.key, newKey)) {
+            counter.decrement(holder);
+            if (holder.count == 0) {
+                result.get(counter.key).remove(holder.value);
             }
-            result.computeIfAbsent(oldKey, k -> setSupplier.get()).add(newValue);
+            lastHolder = counter.add(newValue);
+            if (lastHolder.count == 1) {
+                result.computeIfAbsent(counter.key, k -> setSupplier.get()).add(newValue);
+            }
+            lastCounter = counter;
         } else {
-            remove(oldKey, oldValue);
+            remove(counter, holder);
             add(newKey, newValue);
         }
     }
 
     @Override
-    public void remove(Key_ key, Value_ value) {
-        var counter = valueCounts.get(key);
-        var newCount = counter.remove(value);
-        if (newCount == 0) {
-            result.get(key).remove(value);
+    public void remove(ToMapPerKeyCounter<Key_, Value_> counter, CountHolder<Value_> holder) {
+        counter.decrement(holder);
+        if (holder.count == 0) {
+            result.get(counter.key).remove(holder.value);
         }
         if (counter.isEmpty()) {
-            valueCounts.remove(key);
-            result.remove(key);
+            valueCounts.remove(counter.key);
+            result.remove(counter.key);
         }
     }
 
+    @Override
+    public @Nullable ToMapPerKeyCounter<Key_, Value_> lastCounter() {
+        return lastCounter;
+    }
+
+    @Override
+    public @Nullable CountHolder<Value_> lastHolder() {
+        return lastHolder;
+    }
+
     @Override
     public Result_ getResult() {
         return result;