fix: Revert "perf: group declarative variables after aligned variables when possible (#1725)"

This reverts commit 1a1909c.

Author: Christopher-Chianelli

core/src/main/java/ai/timefold/solver/core/impl/domain/variable/declarative/DefaultShadowVariableSessionFactory.java

@@ -189,14 +189,14 @@ private record GroupVariableUpdaterInfo<Solution_>(
             List<DeclarativeShadowVariableDescriptor<Solution_>> sortedDeclarativeVariableDescriptors,
             List<VariableUpdaterInfo<Solution_>> allUpdaters,
             List<VariableUpdaterInfo<Solution_>> groupedUpdaters,
-            Map<DeclarativeShadowVariableDescriptor<Solution_>, Map<Object, List<VariableUpdaterInfo<Solution_>>>> variableToEntityToGroupUpdater) {
+            Map<DeclarativeShadowVariableDescriptor<Solution_>, Map<Object, VariableUpdaterInfo<Solution_>>> variableToEntityToGroupUpdater) {
 
         public List<VariableUpdaterInfo<Solution_>> getUpdatersForEntityVariable(Object entity,
                 DeclarativeShadowVariableDescriptor<Solution_> declarativeShadowVariableDescriptor) {
             if (variableToEntityToGroupUpdater.containsKey(declarativeShadowVariableDescriptor)) {
-                var updaters = variableToEntityToGroupUpdater.get(declarativeShadowVariableDescriptor).get(entity);
-                if (updaters != null) {
-                    return updaters;
+                var updater = variableToEntityToGroupUpdater.get(declarativeShadowVariableDescriptor).get(entity);
+                if (updater != null) {
+                    return List.of(updater);
                 }
             }
             for (var shadowVariableDescriptor : sortedDeclarativeVariableDescriptors) {
@@ -225,7 +225,24 @@ public List<VariableUpdaterInfo<Solution_>> getUpdatersForEntityVariable(Object
         var groupVariables = new ArrayList<DeclarativeShadowVariableDescriptor<Solution_>>();
         groupIndexToVariables.put(0, groupVariables);
         for (var declarativeShadowVariableDescriptor : sortedDeclarativeVariableDescriptors) {
-            if (shouldCreateNewGroupForVariable(declarativeShadowVariableDescriptor, groupVariables)) {
+            // If a @ShadowSources has a group source (i.e. "visitGroup[].arrivalTimes"),
+            // create a new group since it must wait until all members of that group are processed
+            var hasGroupSources = Arrays.stream(declarativeShadowVariableDescriptor.getSources())
+                    .anyMatch(rootVariableSource -> rootVariableSource.parentVariableType() == ParentVariableType.GROUP);
+
+            // If a @ShadowSources has an alignment key,
+            // create a new group since multiple entities must be updated for this node
+            var hasAlignmentKey = declarativeShadowVariableDescriptor.getAlignmentKeyMap() != null;
+
+            // If the previous @ShadowSources has an alignment key,
+            // create a new group since we are updating a single entity again
+            // NOTE: Can potentially be optimized/share a node if VariableUpdaterInfo
+            //       update each group member independently after the alignmentKey
+            var previousHasAlignmentKey = !groupVariables.isEmpty() && groupVariables.get(0).getAlignmentKeyMap() != null;
+
+            if (!groupVariables.isEmpty() && (hasGroupSources
+                    || hasAlignmentKey
+                    || previousHasAlignmentKey)) {
                 groupVariables = new ArrayList<>();
                 groupIndexToVariables.put(groupIndexToVariables.size(), groupVariables);
             }
@@ -235,12 +252,11 @@ public List<VariableUpdaterInfo<Solution_>> getUpdatersForEntityVariable(Object
         var out = new HashMap<VariableMetaModel<Solution_, ?, ?>, GroupVariableUpdaterInfo<Solution_>>();
         var allUpdaters = new ArrayList<VariableUpdaterInfo<Solution_>>();
         var groupedUpdaters =
-                new HashMap<DeclarativeShadowVariableDescriptor<Solution_>, Map<Object, List<VariableUpdaterInfo<Solution_>>>>();
+                new HashMap<DeclarativeShadowVariableDescriptor<Solution_>, Map<Object, VariableUpdaterInfo<Solution_>>>();
         var updaterKey = 0;
         for (var entryKey = 0; entryKey < groupIndexToVariables.size(); entryKey++) {
             var entryGroupVariables = groupIndexToVariables.get(entryKey);
             var updaters = new ArrayList<VariableUpdaterInfo<Solution_>>();
-            var alignmentKeyToGroupIndex = new HashMap<Object, Integer>();
             for (var declarativeShadowVariableDescriptor : entryGroupVariables) {
                 var updater = new VariableUpdaterInfo<>(
                         declarativeShadowVariableDescriptor.getVariableMetaModel(),
@@ -249,11 +265,31 @@ public List<VariableUpdaterInfo<Solution_>> getUpdatersForEntityVariable(Object
                         declarativeShadowVariableDescriptor.getEntityDescriptor().getShadowVariableLoopedDescriptor(),
                         declarativeShadowVariableDescriptor.getMemberAccessor(),
                         declarativeShadowVariableDescriptor.getCalculator()::executeGetter);
-                if (entryGroupVariables.get(0).getAlignmentKeyMap() != null) {
-                    updaterKey = processAlignmentGroupVariableAndGetNextUpdaterKey(entities,
-                            declarativeShadowVariableDescriptor, entryGroupVariables, updater,
-                            updaterKey, updaters,
-                            allUpdaters, alignmentKeyToGroupIndex, groupedUpdaters);
+                if (declarativeShadowVariableDescriptor.getAlignmentKeyMap() != null) {
+                    var alignmentKeyFunction = declarativeShadowVariableDescriptor.getAlignmentKeyMap();
+                    var alignmentKeyToAlignedEntitiesMap = new HashMap<Object, List<Object>>();
+                    for (var entity : entities) {
+                        if (declarativeShadowVariableDescriptor.getEntityDescriptor().getEntityClass().isInstance(entity)) {
+                            var alignmentKey = alignmentKeyFunction.apply(entity);
+                            alignmentKeyToAlignedEntitiesMap.computeIfAbsent(alignmentKey, k -> new ArrayList<>()).add(entity);
+                        }
+                    }
+                    for (var alignmentGroup : alignmentKeyToAlignedEntitiesMap.entrySet()) {
+                        var updaterCopy = updater.withGroupId(updaterKey);
+                        if (alignmentGroup.getKey() == null) {
+                            updaters.add(updaterCopy);
+                            allUpdaters.add(updaterCopy);
+                        } else {
+                            updaterCopy = updaterCopy.withGroupEntities(alignmentGroup.getValue().toArray(new Object[0]));
+                            var variableUpdaterMap = groupedUpdaters.computeIfAbsent(declarativeShadowVariableDescriptor,
+                                    ignored -> new IdentityHashMap<>());
+                            for (var entity : alignmentGroup.getValue()) {
+                                variableUpdaterMap.put(entity, updaterCopy);
+                            }
+                        }
+                        updaterKey++;
+                    }
+                    updaterKey--; // it will be incremented again at end of the loop
                 } else {
                     updaters.add(updater);
                     allUpdaters.add(updater);
@@ -265,72 +301,12 @@ public List<VariableUpdaterInfo<Solution_>> getUpdatersForEntityVariable(Object
             for (var declarativeShadowVariableDescriptor : entryGroupVariables) {
                 out.put(declarativeShadowVariableDescriptor.getVariableMetaModel(), groupVariableUpdaterInfo);
             }
-            if (entryGroupVariables.get(0).getAlignmentKeyMap() == null) {
-                updaterKey++;
-            }
+            updaterKey++;
         }
         allUpdaters.replaceAll(updater -> updater.withGroupId(groupIndexToVariables.size()));
         return out;
     }
 
-    private static <Solution_> boolean shouldCreateNewGroupForVariable(
-            DeclarativeShadowVariableDescriptor<Solution_> declarativeShadowVariableDescriptor,
-            List<DeclarativeShadowVariableDescriptor<Solution_>> groupVariables) {
-        // If a @ShadowSources has a group source (i.e. "visitGroup[].arrivalTimes"),
-        // create a new group since it must wait until all members of that group are processed
-        var hasGroupSources = Arrays.stream(declarativeShadowVariableDescriptor.getSources())
-                .anyMatch(rootVariableSource -> rootVariableSource.parentVariableType() == ParentVariableType.GROUP);
-
-        // If a @ShadowSources has an alignment key,
-        // create a new group since multiple entities must be updated for this node
-        var alignmentKey = declarativeShadowVariableDescriptor.getAlignmentKeyName();
-        var previousAlignmentKey = groupVariables.isEmpty() ? null : groupVariables.get(0).getAlignmentKeyName();
-
-        return !groupVariables.isEmpty() && (hasGroupSources
-                || (alignmentKey != null && !Objects.equals(alignmentKey, previousAlignmentKey)));
-    }
-
-    private static <Solution_> int processAlignmentGroupVariableAndGetNextUpdaterKey(Object[] entities,
-            DeclarativeShadowVariableDescriptor<Solution_> declarativeShadowVariableDescriptor,
-            List<DeclarativeShadowVariableDescriptor<Solution_>> entryGroupVariables, VariableUpdaterInfo<Solution_> updater,
-            int updaterKey, List<VariableUpdaterInfo<Solution_>> updaters,
-            List<VariableUpdaterInfo<Solution_>> allUpdaters, HashMap<Object, Integer> alignmentKeyToGroupIndex,
-            Map<DeclarativeShadowVariableDescriptor<Solution_>, Map<Object, List<VariableUpdaterInfo<Solution_>>>> groupedUpdaters) {
-        var alignmentKeyFunction = entryGroupVariables.get(0).getAlignmentKeyMap();
-        var alignmentKeyToAlignedEntitiesMap = new HashMap<Object, List<Object>>();
-        for (var entity : entities) {
-            if (declarativeShadowVariableDescriptor.getEntityDescriptor().getEntityClass().isInstance(entity)) {
-                var alignmentKey = alignmentKeyFunction.apply(entity);
-                alignmentKeyToAlignedEntitiesMap.computeIfAbsent(alignmentKey, k -> new ArrayList<>()).add(entity);
-            }
-        }
-        for (var alignmentGroup : alignmentKeyToAlignedEntitiesMap.entrySet()) {
-            if (alignmentGroup.getKey() == null) {
-                var updaterCopy = updater.withGroupId(updaterKey);
-                updaters.add(updaterCopy);
-                allUpdaters.add(updaterCopy);
-                updaterKey++;
-            } else {
-                final var newAlignmentUpdaterKey = updaterKey;
-                final var alignmentUpdaterKey = (int) alignmentKeyToGroupIndex.computeIfAbsent(alignmentGroup.getKey(),
-                        ignored -> newAlignmentUpdaterKey);
-                if (alignmentUpdaterKey == newAlignmentUpdaterKey) {
-                    updaterKey++;
-                }
-                var updaterCopy = updater.withGroupId(alignmentUpdaterKey);
-                updaterCopy = updaterCopy.withGroupEntities(alignmentGroup.getValue().toArray(new Object[0]),
-                        declarativeShadowVariableDescriptor.getAlignmentKeyName() != null);
-                var variableUpdaterMap = groupedUpdaters.computeIfAbsent(declarativeShadowVariableDescriptor,
-                        ignored -> groupedUpdaters.getOrDefault(entryGroupVariables.get(0),
-                                new IdentityHashMap<>()));
-                for (var entity : alignmentGroup.getValue()) {
-                    variableUpdaterMap.computeIfAbsent(entity, ignored -> new ArrayList<>()).add(updaterCopy);
-                }
-            }
-        }
-        return updaterKey;
-    }
-
     private static <Solution_> VariableReferenceGraph buildArbitrarySingleEntityGraph(
             SolutionDescriptor<Solution_> solutionDescriptor,
             VariableReferenceGraphBuilder<Solution_> variableReferenceGraphBuilder, Object[] entities,

core/src/main/java/ai/timefold/solver/core/impl/domain/variable/declarative/DefaultVariableReferenceGraph.java

@@ -19,16 +19,9 @@ public DefaultVariableReferenceGraph(VariableReferenceGraphBuilder<Solution_> ou
 
         var entityToVariableReferenceMap = new IdentityHashMap<Object, List<GraphNode<Solution_>>>();
         for (var instance : nodeList) {
-            if (instance.groupEntityIds() == null) {
-                var entity = instance.entity();
-                entityToVariableReferenceMap.computeIfAbsent(entity, ignored -> new ArrayList<>())
-                        .add(instance);
-            } else {
-                for (var groupEntity : instance.variableReferences().get(0).groupEntities()) {
-                    entityToVariableReferenceMap.computeIfAbsent(groupEntity, ignored -> new ArrayList<>())
-                            .add(instance);
-                }
-            }
+            var entity = instance.entity();
+            entityToVariableReferenceMap.computeIfAbsent(entity, ignored -> new ArrayList<>())
+                    .add(instance);
         }
         // Immutable optimized version of the map, now that it won't be updated anymore.
         var immutableEntityToVariableReferenceMap = mapOfListsDeepCopyOf(entityToVariableReferenceMap);

core/src/main/java/ai/timefold/solver/core/impl/domain/variable/declarative/VariableUpdaterInfo.java

@@ -18,48 +18,29 @@ public record VariableUpdaterInfo<Solution_>(
         @Nullable ShadowVariableLoopedVariableDescriptor<Solution_> shadowVariableLoopedDescriptor,
         MemberAccessor memberAccessor,
         Function<Object, Object> calculator,
-        @Nullable Object[] groupEntities,
-        boolean isGroupAligned) {
+        @Nullable Object[] groupEntities) {
 
     public VariableUpdaterInfo(VariableMetaModel<Solution_, ?, ?> id,
             int groupId,
             DeclarativeShadowVariableDescriptor<Solution_> variableDescriptor,
             @Nullable ShadowVariableLoopedVariableDescriptor<Solution_> shadowVariableLoopedDescriptor,
             MemberAccessor memberAccessor,
             Function<Object, Object> calculator) {
-        // isGroupAligned defaults to true, so we can just check it instead of checking
-        // if groupEntities is null before determining what updateIfChanged to call
-        this(id, groupId, variableDescriptor, shadowVariableLoopedDescriptor, memberAccessor, calculator, null, true);
+        this(id, groupId, variableDescriptor, shadowVariableLoopedDescriptor, memberAccessor, calculator, null);
     }
 
     public VariableUpdaterInfo<Solution_> withGroupId(int groupId) {
         return new VariableUpdaterInfo<>(id, groupId, variableDescriptor, shadowVariableLoopedDescriptor, memberAccessor,
-                calculator, groupEntities, isGroupAligned);
+                calculator, groupEntities);
     }
 
-    public VariableUpdaterInfo<Solution_> withGroupEntities(Object[] groupEntities, boolean isGroupAligned) {
+    public VariableUpdaterInfo<Solution_> withGroupEntities(Object[] groupEntities) {
         return new VariableUpdaterInfo<>(id, groupId, variableDescriptor, shadowVariableLoopedDescriptor, memberAccessor,
-                calculator, groupEntities, isGroupAligned);
+                calculator, groupEntities);
     }
 
     public boolean updateIfChanged(Object entity, ChangedVariableNotifier<Solution_> changedVariableNotifier) {
-        if (isGroupAligned) {
-            return updateIfChanged(entity, calculator.apply(entity), changedVariableNotifier);
-        } else {
-            var anyChanged = false;
-            for (var groupEntity : groupEntities) {
-                var oldValue = variableDescriptor.getValue(groupEntity);
-                var newValue = calculator.apply(groupEntity);
-
-                if (!Objects.equals(oldValue, newValue)) {
-                    changedVariableNotifier.beforeVariableChanged().accept(variableDescriptor, groupEntity);
-                    variableDescriptor.setValue(groupEntity, newValue);
-                    changedVariableNotifier.afterVariableChanged().accept(variableDescriptor, groupEntity);
-                    anyChanged = true;
-                }
-            }
-            return anyChanged;
-        }
+        return updateIfChanged(entity, calculator.apply(entity), changedVariableNotifier);
     }
 
     public boolean updateIfChanged(Object entity, @Nullable Object newValue,

core/src/test/java/ai/timefold/solver/core/impl/domain/variable/listener/support/VariableListenerSupportTest.java

@@ -297,15 +297,12 @@ void shadowVariableListGraphEvents() {
                 // which is the first member of the group
                 for (var element : visit.getConcurrentValueGroup()) {
                     verifyAddEdge.accept(serviceReadyTime, element, serviceStartTime, visit);
-                    // start and finish time use the same node, so no edge between them
+                    verifyAddEdge.accept(serviceStartTime, visit, serviceFinishTime, element);
                 }
             }
 
             if (visit.getPreviousValue() != null) {
-                var previousRepresentative =
-                        visit.getPreviousValue().getConcurrentValueGroup() == null ? visit.getPreviousValue()
-                                : visit.getPreviousValue().getConcurrentValueGroup().get(0);
-                verifyAddEdge.accept(serviceFinishTime, previousRepresentative, serviceReadyTime, visit);
+                verifyAddEdge.accept(serviceFinishTime, visit.getPreviousValue(), serviceReadyTime, visit);
             }
         }
         // Note: addEdge only adds an edge if it does not already exists in the graph,