Skip to content

Commit ec49da1

Browse files
triceotriceo
committed
fix: respect merge function in toMap() collector
1 parent 60afd2f commit ec49da1

File tree

8 files changed

+113
-78
lines changed

8 files changed

+113
-78
lines changed

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,10 @@ public Value_ merge(BinaryOperator<Value_> mergeFunction) {
3535
// The impact is negligible, assuming there are not too many values for the same key.
3636
return counts.entrySet()
3737
.stream()
38-
.flatMap(e -> Stream.generate(e::getKey).limit(e.getValue()))
38+
.map(e -> Stream.generate(e::getKey)
39+
.limit(e.getValue())
40+
.reduce(mergeFunction)
41+
.orElseThrow(() -> new IllegalStateException("Impossible state: Should have had at least one value.")))
3942
.reduce(mergeFunction)
4043
.orElseThrow(() -> new IllegalStateException("Impossible state: Should have had at least one value."));
4144
}

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

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,26 +27,20 @@ public ToSimpleMapResultContainer(IntFunction<Result_> resultSupplier, BinaryOpe
2727
@Override
2828
public void add(Key_ key, Value_ value) {
2929
ToMapPerKeyCounter<Value_> counter = valueCounts.computeIfAbsent(key, k -> new ToMapPerKeyCounter<>());
30-
long newCount = counter.add(value);
31-
if (newCount == 1L) {
32-
result.put(key, value);
33-
} else {
34-
result.put(key, counter.merge(mergeFunction));
35-
}
30+
counter.add(value);
31+
result.put(key, counter.merge(mergeFunction));
3632
}
3733

3834
@Override
3935
public void remove(Key_ key, Value_ value) {
4036
ToMapPerKeyCounter<Value_> counter = valueCounts.get(key);
41-
long newCount = counter.remove(value);
42-
if (newCount == 0L) {
37+
counter.remove(value);
38+
if (counter.isEmpty()) {
4339
result.remove(key);
40+
valueCounts.remove(key);
4441
} else {
4542
result.put(key, counter.merge(mergeFunction));
4643
}
47-
if (counter.isEmpty()) {
48-
valueCounts.remove(key);
49-
}
5044
}
5145

5246
@Override

‎core/src/test/java/ai/timefold/solver/core/impl/score/stream/bavet/BavetAdvancedGroupByConstraintStreamTest.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,14 @@
1212
import static ai.timefold.solver.core.impl.testdata.util.PlannerTestUtils.asMap;
1313
import static org.assertj.core.api.Assertions.assertThatCode;
1414

15+
import java.util.Arrays;
1516
import java.util.Objects;
1617
import java.util.Set;
1718
import java.util.function.Function;
1819
import java.util.stream.Stream;
1920

2021
import ai.timefold.solver.core.api.score.buildin.simple.SimpleScore;
22+
import ai.timefold.solver.core.api.score.stream.ConstraintCollectors;
2123
import ai.timefold.solver.core.api.score.stream.Joiners;
2224
import ai.timefold.solver.core.impl.score.director.InnerScoreDirector;
2325
import ai.timefold.solver.core.impl.score.stream.common.AbstractAdvancedGroupByConstraintStreamTest;
@@ -152,15 +154,15 @@ void biGroupByRecollected() {
152154
.forEachUniquePair(TestdataLavishEntity.class, equal(TestdataLavishEntity::getEntityGroup))
153155
// Stream of all unique entity bi tuples that share a group
154156
.groupBy((a, b) -> a.getEntityGroup(), countBi())
155-
.groupBy(toMap((g, c) -> g, (g, c) -> c, Integer::sum))
157+
.groupBy(ConstraintCollectors.toList((a, b) -> a))
156158
.penalize(SimpleScore.ONE)
157159
.asConstraint(TEST_CONSTRAINT_NAME));
158160

159161
// From scratch
160162
scoreDirector.setWorkingSolution(solution);
161163
assertScore(scoreDirector,
162164
assertMatchWithScore(-1,
163-
asMap(solution.getFirstEntityGroup(), 3, solution.getEntityGroupList().get(1), 1)));
165+
Arrays.asList(solution.getFirstEntityGroup(), solution.getEntityGroupList().get(1))));
164166

165167
// Incremental
166168
TestdataLavishEntity entity = solution.getFirstEntity();
@@ -169,7 +171,7 @@ void biGroupByRecollected() {
169171
scoreDirector.afterEntityRemoved(entity);
170172
assertScore(scoreDirector,
171173
assertMatchWithScore(-1,
172-
asMap(solution.getFirstEntityGroup(), 1, solution.getEntityGroupList().get(1), 1)));
174+
Arrays.asList(solution.getEntityGroupList().get(1), solution.getFirstEntityGroup())));
173175
}
174176

175177
@TestTemplate

‎core/src/test/java/ai/timefold/solver/core/impl/score/stream/collector/bi/InnerBiConstraintCollectorsTest.java‎

Lines changed: 19 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -816,29 +816,29 @@ public void toMapDuplicates() { // PLANNER-2271
816816
@Override
817817
@Test
818818
public void toMapMerged() {
819-
BiConstraintCollector<Integer, Integer, ?, Map<Integer, Integer>> collector = ConstraintCollectors.toMap(Integer::sum,
820-
Integer::sum, Integer::sum);
819+
BiConstraintCollector<Integer, Integer, ?, Map<Integer, Integer>> collector =
820+
ConstraintCollectors.toMap((a, b) -> a, (a, b) -> b, Integer::sum);
821821
Object container = collector.supplier().get();
822822

823823
// Default state.
824824
assertResult(collector, container, emptyMap());
825825
// Add first value, we have one now.
826826
int firstValue = 2;
827-
Runnable firstRetractor = accumulate(collector, container, firstValue, 0);
828-
assertResult(collector, container, asMap(2, 2));
827+
Runnable firstRetractor = accumulate(collector, container, firstValue, 1);
828+
assertResult(collector, container, asMap(2, 1));
829829
// Add second value, we have two now.
830830
int secondValue = 1;
831-
Runnable secondRetractor = accumulate(collector, container, secondValue, 0);
832-
assertResult(collector, container, asMap(2, 2, 1, 1));
831+
Runnable secondRetractor = accumulate(collector, container, secondValue, 1);
832+
assertResult(collector, container, asMap(2, 1, 1, 1));
833833
// Add third value, same as the second. We now have three values, two of which map to the same key.
834-
Runnable thirdRetractor = accumulate(collector, container, secondValue, 0);
835-
assertResult(collector, container, asMap(2, 2, 1, 2));
834+
Runnable thirdRetractor = accumulate(collector, container, secondValue, 2);
835+
assertResult(collector, container, asMap(2, 1, 1, 3));
836836
// Retract one instance of the second value; we only have two values now.
837837
secondRetractor.run();
838-
assertResult(collector, container, asMap(2, 2, 1, 1));
838+
assertResult(collector, container, asMap(2, 1, 1, 2));
839839
// Retract final instance of the second value; we only have one value now.
840840
thirdRetractor.run();
841-
assertResult(collector, container, asMap(2, 2));
841+
assertResult(collector, container, asMap(2, 1));
842842
// Retract last value; there are no values now.
843843
firstRetractor.run();
844844
assertResult(collector, container, emptyMap());
@@ -879,28 +879,28 @@ public void toSortedMap() {
879879
@Test
880880
public void toSortedMapMerged() {
881881
BiConstraintCollector<Integer, Integer, ?, SortedMap<Integer, Integer>> collector = ConstraintCollectors
882-
.toSortedMap(Integer::sum, Integer::sum, Integer::sum);
882+
.toSortedMap((a, b) -> a, (a, b) -> b, Integer::sum);
883883
Object container = collector.supplier().get();
884884

885885
// Default state.
886886
assertResult(collector, container, emptySortedMap());
887887
// Add first value, we have one now.
888888
int firstValue = 2;
889-
Runnable firstRetractor = accumulate(collector, container, firstValue, 0);
890-
assertResult(collector, container, asSortedMap(2, 2));
889+
Runnable firstRetractor = accumulate(collector, container, firstValue, 1);
890+
assertResult(collector, container, asSortedMap(2, 1));
891891
// Add second value, we have two now.
892892
int secondValue = 1;
893-
Runnable secondRetractor = accumulate(collector, container, secondValue, 0);
894-
assertResult(collector, container, asSortedMap(2, 2, 1, 1));
893+
Runnable secondRetractor = accumulate(collector, container, secondValue, 1);
894+
assertResult(collector, container, asSortedMap(2, 1, 1, 1));
895895
// Add third value, same as the second. We now have three values, two of which map to the same key.
896-
Runnable thirdRetractor = accumulate(collector, container, secondValue, 0);
897-
assertResult(collector, container, asSortedMap(2, 2, 1, 2));
896+
Runnable thirdRetractor = accumulate(collector, container, secondValue, 2);
897+
assertResult(collector, container, asSortedMap(2, 1, 1, 3));
898898
// Retract one instance of the second value; we only have two values now.
899899
secondRetractor.run();
900-
assertResult(collector, container, asSortedMap(2, 2, 1, 1));
900+
assertResult(collector, container, asSortedMap(2, 1, 1, 2));
901901
// Retract final instance of the second value; we only have one value now.
902902
thirdRetractor.run();
903-
assertResult(collector, container, asSortedMap(2, 2));
903+
assertResult(collector, container, asSortedMap(2, 1));
904904
// Retract last value; there are no values now.
905905
firstRetractor.run();
906906
assertResult(collector, container, emptySortedMap());

‎core/src/test/java/ai/timefold/solver/core/impl/score/stream/collector/quad/InnerQuadConstraintCollectorsTest.java‎

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -867,28 +867,28 @@ public void toMap() {
867867
@Test
868868
public void toMapMerged() {
869869
QuadConstraintCollector<Integer, Integer, Integer, Integer, ?, Map<Integer, Integer>> collector = ConstraintCollectors
870-
.toMap((a, b, c, d) -> a + b + c + d, (a, b, c, d) -> a + b + c + d, Integer::sum);
870+
.toMap((a, b, c, d) -> a + b, (a, b, c, d) -> c + d, Integer::sum);
871871
Object container = collector.supplier().get();
872872

873873
// Default state.
874874
assertResult(collector, container, emptyMap());
875875
// Add first value, we have one now.
876876
int firstValue = 2;
877-
Runnable firstRetractor = accumulate(collector, container, firstValue, 0, 0, 0);
878-
assertResult(collector, container, asMap(2, 2));
877+
Runnable firstRetractor = accumulate(collector, container, firstValue, 1, 3, 2);
878+
assertResult(collector, container, asMap(3, 5));
879879
// Add second value, we have two now.
880880
int secondValue = 1;
881-
Runnable secondRetractor = accumulate(collector, container, secondValue, 0, 0, 0);
882-
assertResult(collector, container, asMap(2, 2, 1, 1));
881+
Runnable secondRetractor = accumulate(collector, container, secondValue, 3, 3, 2);
882+
assertResult(collector, container, asMap(3, 5, 4, 5));
883883
// Add third value, same as the second. We now have three values, two of which map to the same key.
884-
Runnable thirdRetractor = accumulate(collector, container, secondValue, 0, 0, 0);
885-
assertResult(collector, container, asMap(2, 2, 1, 2));
884+
Runnable thirdRetractor = accumulate(collector, container, secondValue, 3, 3, 1);
885+
assertResult(collector, container, asMap(3, 5, 4, 9));
886886
// Retract one instance of the second value; we only have two values now.
887887
secondRetractor.run();
888-
assertResult(collector, container, asMap(2, 2, 1, 1));
888+
assertResult(collector, container, asMap(3, 5, 4, 4));
889889
// Retract final instance of the second value; we only have one value now.
890890
thirdRetractor.run();
891-
assertResult(collector, container, asMap(2, 2));
891+
assertResult(collector, container, asMap(3, 5));
892892
// Retract last value; there are no values now.
893893
firstRetractor.run();
894894
assertResult(collector, container, emptyMap());
@@ -929,28 +929,29 @@ public void toSortedMap() {
929929
@Test
930930
public void toSortedMapMerged() {
931931
QuadConstraintCollector<Integer, Integer, Integer, Integer, ?, SortedMap<Integer, Integer>> collector =
932-
ConstraintCollectors.toSortedMap((a, b, c, d) -> a + b + c + d, (a, b, c, d) -> a + b + c + d, Integer::sum);
932+
ConstraintCollectors
933+
.toSortedMap((a, b, c, d) -> a + b, (a, b, c, d) -> c + d, Integer::sum);
933934
Object container = collector.supplier().get();
934935

935936
// Default state.
936937
assertResult(collector, container, emptySortedMap());
937938
// Add first value, we have one now.
938939
int firstValue = 2;
939-
Runnable firstRetractor = accumulate(collector, container, firstValue, 0, 0, 0);
940-
assertResult(collector, container, asSortedMap(2, 2));
940+
Runnable firstRetractor = accumulate(collector, container, firstValue, 1, 3, 2);
941+
assertResult(collector, container, asSortedMap(3, 5));
941942
// Add second value, we have two now.
942943
int secondValue = 1;
943-
Runnable secondRetractor = accumulate(collector, container, secondValue, 0, 0, 0);
944-
assertResult(collector, container, asSortedMap(2, 2, 1, 1));
944+
Runnable secondRetractor = accumulate(collector, container, secondValue, 3, 3, 2);
945+
assertResult(collector, container, asSortedMap(3, 5, 4, 5));
945946
// Add third value, same as the second. We now have three values, two of which map to the same key.
946-
Runnable thirdRetractor = accumulate(collector, container, secondValue, 0, 0, 0);
947-
assertResult(collector, container, asSortedMap(2, 2, 1, 2));
947+
Runnable thirdRetractor = accumulate(collector, container, secondValue, 3, 3, 1);
948+
assertResult(collector, container, asSortedMap(3, 5, 4, 9));
948949
// Retract one instance of the second value; we only have two values now.
949950
secondRetractor.run();
950-
assertResult(collector, container, asSortedMap(2, 2, 1, 1));
951+
assertResult(collector, container, asSortedMap(3, 5, 4, 4));
951952
// Retract final instance of the second value; we only have one value now.
952953
thirdRetractor.run();
953-
assertResult(collector, container, asSortedMap(2, 2));
954+
assertResult(collector, container, asSortedMap(3, 5));
954955
// Retract last value; there are no values now.
955956
firstRetractor.run();
956957
assertResult(collector, container, emptySortedMap());

‎core/src/test/java/ai/timefold/solver/core/impl/score/stream/collector/tri/InnerTriConstraintCollectorsTest.java‎

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -820,28 +820,28 @@ public void toMap() {
820820
@Test
821821
public void toMapMerged() {
822822
TriConstraintCollector<Integer, Integer, Integer, ?, Map<Integer, Integer>> collector = ConstraintCollectors
823-
.toMap((a, b, c) -> a + b + c, (a, b, c) -> a + b + c, Integer::sum);
823+
.toMap((a, b, c) -> a, (a, b, c) -> b + c, Integer::sum);
824824
Object container = collector.supplier().get();
825825

826826
// Default state.
827827
assertResult(collector, container, emptyMap());
828828
// Add first value, we have one now.
829829
int firstValue = 2;
830-
Runnable firstRetractor = accumulate(collector, container, firstValue, 0, 0);
831-
assertResult(collector, container, asMap(2, 2));
830+
Runnable firstRetractor = accumulate(collector, container, firstValue, 1, 2);
831+
assertResult(collector, container, asMap(2, 3));
832832
// Add second value, we have two now.
833833
int secondValue = 1;
834-
Runnable secondRetractor = accumulate(collector, container, secondValue, 0, 0);
835-
assertResult(collector, container, asMap(2, 2, 1, 1));
834+
Runnable secondRetractor = accumulate(collector, container, secondValue, 1, 2);
835+
assertResult(collector, container, asMap(2, 3, 1, 3));
836836
// Add third value, same as the second. We now have three values, two of which map to the same key.
837-
Runnable thirdRetractor = accumulate(collector, container, secondValue, 0, 0);
838-
assertResult(collector, container, asMap(2, 2, 1, 2));
837+
Runnable thirdRetractor = accumulate(collector, container, secondValue, 2, 3);
838+
assertResult(collector, container, asMap(2, 3, 1, 8));
839839
// Retract one instance of the second value; we only have two values now.
840840
secondRetractor.run();
841-
assertResult(collector, container, asMap(2, 2, 1, 1));
841+
assertResult(collector, container, asMap(2, 3, 1, 5));
842842
// Retract final instance of the second value; we only have one value now.
843843
thirdRetractor.run();
844-
assertResult(collector, container, asMap(2, 2));
844+
assertResult(collector, container, asMap(2, 3));
845845
// Retract last value; there are no values now.
846846
firstRetractor.run();
847847
assertResult(collector, container, emptyMap());
@@ -882,28 +882,28 @@ public void toSortedMap() {
882882
@Test
883883
public void toSortedMapMerged() {
884884
TriConstraintCollector<Integer, Integer, Integer, ?, SortedMap<Integer, Integer>> collector = ConstraintCollectors
885-
.toSortedMap((a, b, c) -> a + b + c, (a, b, c) -> a + b + c, Integer::sum);
885+
.toSortedMap((a, b, c) -> a, (a, b, c) -> b + c, Integer::sum);
886886
Object container = collector.supplier().get();
887887

888888
// Default state.
889889
assertResult(collector, container, emptySortedMap());
890890
// Add first value, we have one now.
891891
int firstValue = 2;
892-
Runnable firstRetractor = accumulate(collector, container, firstValue, 0, 0);
893-
assertResult(collector, container, asSortedMap(2, 2));
892+
Runnable firstRetractor = accumulate(collector, container, firstValue, 1, 2);
893+
assertResult(collector, container, asSortedMap(2, 3));
894894
// Add second value, we have two now.
895895
int secondValue = 1;
896-
Runnable secondRetractor = accumulate(collector, container, secondValue, 0, 0);
897-
assertResult(collector, container, asSortedMap(2, 2, 1, 1));
896+
Runnable secondRetractor = accumulate(collector, container, secondValue, 1, 2);
897+
assertResult(collector, container, asSortedMap(2, 3, 1, 3));
898898
// Add third value, same as the second. We now have three values, two of which map to the same key.
899-
Runnable thirdRetractor = accumulate(collector, container, secondValue, 0, 0);
900-
assertResult(collector, container, asSortedMap(2, 2, 1, 2));
899+
Runnable thirdRetractor = accumulate(collector, container, secondValue, 2, 3);
900+
assertResult(collector, container, asSortedMap(2, 3, 1, 8));
901901
// Retract one instance of the second value; we only have two values now.
902902
secondRetractor.run();
903-
assertResult(collector, container, asSortedMap(2, 2, 1, 1));
903+
assertResult(collector, container, asSortedMap(2, 3, 1, 5));
904904
// Retract final instance of the second value; we only have one value now.
905905
thirdRetractor.run();
906-
assertResult(collector, container, asSortedMap(2, 2));
906+
assertResult(collector, container, asSortedMap(2, 3));
907907
// Retract last value; there are no values now.
908908
firstRetractor.run();
909909
assertResult(collector, container, emptySortedMap());

0 commit comments

Comments
 (0)