perf: Use a single node to represent multiple shadow variables in graphs (#1673)

When there is only a single declarative entity and all non-declarative
parents go in the same direction (ex: previous), then some variables can
share the same node in the graph. In particular, there is a static
topological order when only the variables on a single entity is
considered, and as long as the next declarative shadow variable does not
have any new non-declarative parents (i.e. group), it can be grouped
into the same node as the previous one. All nodes share the same node if
no groups are used.

Also fixed a bug in the topological sort done to determine variable
trigger order when groups are used (previously it added group variables,
which caused it to create a cycle and make all topological orders
valid).

---------

Co-authored-by: Lukáš Petrovický <[email protected]>

Author: Christopher-Chianelli

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

@@ -51,7 +51,7 @@ public void accept(BitSet changed) {
             }
             visited.set(nextNode);
             var shadowVariable = instanceList.get(nextNode);
-            var isChanged = updateShadowVariable(shadowVariable, graph.isLooped(loopedTracker, nextNode));
+            var isChanged = updateEntityShadowVariables(shadowVariable, graph.isLooped(loopedTracker, nextNode));
 
             if (isChanged) {
                 var iterator = graph.nodeForwardEdges(nextNode);
@@ -96,7 +96,7 @@ private void updateLoopedStatusOfAffectedEntity(Object affectedEntity) {
         for (var node : entityVariablePairFunction.apply(affectedEntity)) {
             // All variables come from the same entity,
             // therefore all have the same looped marker.
-            shadowVariableLoopedDescriptor = node.variableReference().shadowVariableLoopedDescriptor();
+            shadowVariableLoopedDescriptor = node.variableReferences().get(0).shadowVariableLoopedDescriptor();
             if (graph.isLooped(loopedTracker, node.graphNodeId())) {
                 isEntityLooped = true;
                 break;
@@ -114,19 +114,31 @@ private void updateLoopedStatusOfAffectedEntity(Object affectedEntity) {
 
     }
 
-    private boolean updateShadowVariable(EntityVariablePair<Solution_> entityVariable, boolean isLooped) {
+    private boolean updateEntityShadowVariables(EntityVariablePair<Solution_> entityVariable, boolean isLooped) {
         var entity = entityVariable.entity();
-        var shadowVariableReference = entityVariable.variableReference();
-        var oldValue = shadowVariableReference.memberAccessor().executeGetter(entity);
-        var loopDescriptor = shadowVariableReference.shadowVariableLoopedDescriptor();
+        var shadowVariableReferences = entityVariable.variableReferences();
+        var loopDescriptor = shadowVariableReferences.get(0).shadowVariableLoopedDescriptor();
+        var anyChanged = false;
+
         if (loopDescriptor != null) {
-            var oldLooped = (boolean) loopDescriptor.getValue(entity);
-            if (oldLooped != isLooped) {
+            var oldLooped = loopDescriptor.getValue(entity);
+            if (!Objects.equals(oldLooped, isLooped)) {
                 // Loop status change; add to affected entities
                 affectedEntities.add(entityVariable);
+                anyChanged = true;
             }
         }
 
+        for (var shadowVariableReference : shadowVariableReferences) {
+            anyChanged |= updateShadowVariable(entityVariable, isLooped, shadowVariableReference, entity);
+        }
+
+        return anyChanged;
+    }
+
+    private boolean updateShadowVariable(EntityVariablePair<Solution_> entityVariable, boolean isLooped,
+            VariableUpdaterInfo<Solution_> shadowVariableReference, Object entity) {
+        var oldValue = shadowVariableReference.memberAccessor().executeGetter(entity);
         if (isLooped) {
             if (oldValue != null) {
                 affectedEntities.add(entityVariable);
@@ -168,7 +180,7 @@ public AffectedEntities(Consumer<Object> consumer) {
         }
 
         public void add(EntityVariablePair<Solution_> shadowVariable) {
-            var shadowVariableLoopedDescriptor = shadowVariable.variableReference().shadowVariableLoopedDescriptor();
+            var shadowVariableLoopedDescriptor = shadowVariable.variableReferences().get(0).shadowVariableLoopedDescriptor();
             if (shadowVariableLoopedDescriptor == null) {
                 return;
             }

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

@@ -1,6 +1,7 @@
 package ai.timefold.solver.core.impl.domain.variable.declarative;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.BitSet;
 import java.util.Collections;
 import java.util.Comparator;
@@ -15,12 +16,13 @@
 import java.util.Set;
 import java.util.function.IntFunction;
 
+import ai.timefold.solver.core.api.function.TriFunction;
 import ai.timefold.solver.core.impl.domain.solution.descriptor.SolutionDescriptor;
 import ai.timefold.solver.core.impl.domain.variable.descriptor.VariableDescriptor;
 import ai.timefold.solver.core.impl.score.director.InnerScoreDirector;
+import ai.timefold.solver.core.impl.util.MutableInt;
 import ai.timefold.solver.core.preview.api.domain.metamodel.VariableMetaModel;
 
-import org.jspecify.annotations.NonNull;
 import org.jspecify.annotations.NullMarked;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -59,6 +61,8 @@ yield buildSingleDirectionalParentGraph(solutionDescriptor,
                         graphStructureAndDirection,
                         entities);
             }
+            case ARBITRARY_SINGLE_ENTITY_SINGLE_DIRECTIONAL_PARENT_TYPE ->
+                buildArbitrarySingleEntityGraph(solutionDescriptor, variableReferenceGraphBuilder, entities, graphCreator);
             case NO_DYNAMIC_EDGES, ARBITRARY ->
                 buildArbitraryGraph(solutionDescriptor, variableReferenceGraphBuilder, entities, graphCreator);
         };
@@ -81,7 +85,7 @@ static <Solution_> VariableReferenceGraph buildSingleDirectionalParentGraph(
                 topologicalSorter, changedVariableNotifier, entities);
     }
 
-    private static <Solution_> @NonNull List<DeclarativeShadowVariableDescriptor<Solution_>>
+    private static <Solution_> List<DeclarativeShadowVariableDescriptor<Solution_>>
             topologicallySortedDeclarativeShadowVariables(
                     List<DeclarativeShadowVariableDescriptor<Solution_>> declarativeShadowVariables) {
         Map<String, Integer> nameToIndex = new LinkedHashMap<>();
@@ -94,8 +98,10 @@ static <Solution_> VariableReferenceGraph buildSingleDirectionalParentGraph(
             var visited = new HashSet<Integer>();
             for (var source : declarativeShadowVariable.getSources()) {
                 var variableReferences = source.variableSourceReferences();
-                if (variableReferences.size() != 1) {
-                    // variableReferences is from directional variable
+                if (source.parentVariableType() != ParentVariableType.NO_PARENT) {
+                    // We only look at direct usage; if we also added
+                    // edges for groups/directional, we will end up creating a cycle
+                    // which makes all topological orders valid
                     continue;
                 }
                 var variableReference = variableReferences.get(0);
@@ -143,7 +149,7 @@ private static <Solution_> VariableReferenceGraph buildArbitraryGraph(
             VariableReferenceGraphBuilder<Solution_> variableReferenceGraphBuilder, Object[] entities,
             IntFunction<TopologicalOrderGraph> graphCreator) {
         var declarativeShadowVariableDescriptors = solutionDescriptor.getDeclarativeShadowVariableDescriptors();
-        var variableIdToUpdater = new HashMap<VariableMetaModel<?, ?, ?>, VariableUpdaterInfo<Solution_>>();
+        var variableIdToUpdater = EntityVariableUpdaterLookup.<Solution_> entityIndependentLookup();
 
         // Create graph node for each entity/declarative shadow variable pair.
         // Maps a variable id to its source aliases;
@@ -152,7 +158,16 @@ private static <Solution_> VariableReferenceGraph buildArbitraryGraph(
         // to "startTime" of some visit, and thus alias it.
         var declarativeShadowVariableToAliasMap = createGraphNodes(variableReferenceGraphBuilder, entities,
                 declarativeShadowVariableDescriptors, variableIdToUpdater);
+        return buildVariableReferenceGraph(declarativeShadowVariableDescriptors, variableReferenceGraphBuilder,
+                declarativeShadowVariableToAliasMap,
+                graphCreator, entities);
+    }
 
+    private static <Solution_> VariableReferenceGraph buildVariableReferenceGraph(
+            List<DeclarativeShadowVariableDescriptor<Solution_>> declarativeShadowVariableDescriptors,
+            VariableReferenceGraphBuilder<Solution_> variableReferenceGraphBuilder,
+            Map<VariableMetaModel<?, ?, ?>, Set<VariableSourceReference>> declarativeShadowVariableToAliasMap,
+            IntFunction<TopologicalOrderGraph> graphCreator, Object... entities) {
         // Create variable processors for each declarative shadow variable descriptor
         for (var declarativeShadowVariable : declarativeShadowVariableDescriptors) {
             var fromVariableId = declarativeShadowVariable.getVariableMetaModel();
@@ -168,23 +183,130 @@ private static <Solution_> VariableReferenceGraph buildArbitraryGraph(
         return variableReferenceGraphBuilder.build(graphCreator);
     }
 
+    private record GroupVariableUpdaterInfo<Solution_>(
+            List<DeclarativeShadowVariableDescriptor<Solution_>> sortedDeclarativeVariableDescriptors,
+            List<VariableUpdaterInfo<Solution_>> allUpdaters,
+            List<VariableUpdaterInfo<Solution_>> groupedUpdaters) {
+
+        public List<VariableUpdaterInfo<Solution_>> getUpdatersForEntity(Object entity) {
+            for (var shadowVariableDescriptor : sortedDeclarativeVariableDescriptors) {
+                for (var rootSource : shadowVariableDescriptor.getSources()) {
+                    if (rootSource.parentVariableType() == ParentVariableType.GROUP) {
+                        var visitedCount = new MutableInt();
+                        rootSource.valueEntityFunction().accept(entity, ignored -> visitedCount.increment());
+                        if (visitedCount.intValue() > 0) {
+                            return groupedUpdaters;
+                        }
+                    }
+                }
+            }
+            return allUpdaters;
+        }
+
+    }
+
+    private static <Solution_> Map<VariableMetaModel<Solution_, ?, ?>, GroupVariableUpdaterInfo<Solution_>>
+            getGroupVariableUpdaterInfoMap(
+                    List<DeclarativeShadowVariableDescriptor<Solution_>> declarativeShadowVariableDescriptors) {
+        var sortedDeclarativeVariableDescriptors =
+                topologicallySortedDeclarativeShadowVariables(declarativeShadowVariableDescriptors);
+        var groupIndexToVariables = new HashMap<Integer, List<DeclarativeShadowVariableDescriptor<Solution_>>>();
+        var groupVariables = new ArrayList<DeclarativeShadowVariableDescriptor<Solution_>>();
+        groupIndexToVariables.put(0, groupVariables);
+        for (var declarativeShadowVariableDescriptor : sortedDeclarativeVariableDescriptors) {
+            if (!groupVariables.isEmpty() && Arrays.stream(declarativeShadowVariableDescriptor.getSources())
+                    .anyMatch(rootVariableSource -> rootVariableSource.parentVariableType() == ParentVariableType.GROUP)) {
+                // Create a new variable group, since the group might reference prior variables
+                groupVariables = new ArrayList<>();
+                groupIndexToVariables.put(groupIndexToVariables.size(), groupVariables);
+            }
+            groupVariables.add(declarativeShadowVariableDescriptor);
+        }
+
+        var out = new HashMap<VariableMetaModel<Solution_, ?, ?>, GroupVariableUpdaterInfo<Solution_>>();
+        var allUpdaters = new ArrayList<VariableUpdaterInfo<Solution_>>();
+        for (var entryKey = 0; entryKey < groupIndexToVariables.size(); entryKey++) {
+            var entryGroupVariables = groupIndexToVariables.get(entryKey);
+            var updaters = new ArrayList<VariableUpdaterInfo<Solution_>>();
+            for (var declarativeShadowVariableDescriptor : entryGroupVariables) {
+                var updater = new VariableUpdaterInfo<>(
+                        declarativeShadowVariableDescriptor.getVariableMetaModel(),
+                        entryKey,
+                        declarativeShadowVariableDescriptor,
+                        declarativeShadowVariableDescriptor.getEntityDescriptor().getShadowVariableLoopedDescriptor(),
+                        declarativeShadowVariableDescriptor.getMemberAccessor(),
+                        declarativeShadowVariableDescriptor.getCalculator()::executeGetter);
+                updaters.add(updater);
+                allUpdaters.add(updater);
+            }
+            var groupVariableUpdaterInfo =
+                    new GroupVariableUpdaterInfo<Solution_>(sortedDeclarativeVariableDescriptors, allUpdaters, updaters);
+            for (var declarativeShadowVariableDescriptor : entryGroupVariables) {
+                out.put(declarativeShadowVariableDescriptor.getVariableMetaModel(), groupVariableUpdaterInfo);
+            }
+        }
+        allUpdaters.replaceAll(updater -> updater.withGroupId(groupIndexToVariables.size()));
+        return out;
+    }
+
+    private static <Solution_> VariableReferenceGraph buildArbitrarySingleEntityGraph(
+            SolutionDescriptor<Solution_> solutionDescriptor,
+            VariableReferenceGraphBuilder<Solution_> variableReferenceGraphBuilder, Object[] entities,
+            IntFunction<TopologicalOrderGraph> graphCreator) {
+        var declarativeShadowVariableDescriptors = solutionDescriptor.getDeclarativeShadowVariableDescriptors();
+        // Use a dependent lookup; if an entity does not use groups, then all variables can share the same node.
+        // If the entity use groups, then variables must be grouped into their own nodes.
+        var variableIdToUpdater = EntityVariableUpdaterLookup.<Solution_> entityDependentLookup();
+
+        // Create graph node for each entity/declarative shadow variable group pair.
+        // Maps a variable id to the source aliases of all variables in its group;
+        // If the variables are (in topological order)
+        // arrivalTime, readyTime, serviceStartTime, serviceFinishTime,
+        // where serviceStartTime depends on a group of readyTime, then
+        // the groups are [arrivalTime, readyTime] and [serviceStartTime, serviceFinishTime]
+        // this is because from arrivalTime, you can compute readyTime without knowing either
+        // serviceStartTime or serviceFinishTime.
+        var variableIdToGroupedUpdater = getGroupVariableUpdaterInfoMap(declarativeShadowVariableDescriptors);
+        var declarativeShadowVariableToAliasMap = createGraphNodes(variableReferenceGraphBuilder, entities,
+                declarativeShadowVariableDescriptors, variableIdToUpdater,
+                (entity, declarativeShadowVariable, variableId) -> variableIdToGroupedUpdater.get(variableId)
+                        .getUpdatersForEntity(entity));
+        return buildVariableReferenceGraph(declarativeShadowVariableDescriptors, variableReferenceGraphBuilder,
+                declarativeShadowVariableToAliasMap,
+                graphCreator, entities);
+    }
+
+    private static <Solution_> Map<VariableMetaModel<?, ?, ?>, Set<VariableSourceReference>> createGraphNodes(
+            VariableReferenceGraphBuilder<Solution_> graph, Object[] entities,
+            List<DeclarativeShadowVariableDescriptor<Solution_>> declarativeShadowVariableDescriptors,
+            EntityVariableUpdaterLookup<Solution_> variableIdToUpdaters) {
+        return createGraphNodes(graph, entities, declarativeShadowVariableDescriptors, variableIdToUpdaters,
+                (entity, declarativeShadowVariableDescriptor,
+                        variableId) -> Collections.singletonList(new VariableUpdaterInfo<>(
+                                variableId,
+                                variableIdToUpdaters.getNextId(),
+                                declarativeShadowVariableDescriptor,
+                                declarativeShadowVariableDescriptor.getEntityDescriptor().getShadowVariableLoopedDescriptor(),
+                                declarativeShadowVariableDescriptor.getMemberAccessor(),
+                                declarativeShadowVariableDescriptor.getCalculator()::executeGetter)));
+    }
+
     private static <Solution_> Map<VariableMetaModel<?, ?, ?>, Set<VariableSourceReference>> createGraphNodes(
             VariableReferenceGraphBuilder<Solution_> graph, Object[] entities,
             List<DeclarativeShadowVariableDescriptor<Solution_>> declarativeShadowVariableDescriptors,
-            Map<VariableMetaModel<?, ?, ?>, VariableUpdaterInfo<Solution_>> variableIdToUpdater) {
+            EntityVariableUpdaterLookup<Solution_> variableIdToUpdaters,
+            TriFunction<Object, DeclarativeShadowVariableDescriptor<Solution_>, VariableMetaModel<Solution_, ?, ?>, List<VariableUpdaterInfo<Solution_>>> entityVariableToUpdatersMapper) {
         var result = new HashMap<VariableMetaModel<?, ?, ?>, Set<VariableSourceReference>>();
         for (var entity : entities) {
             for (var declarativeShadowVariableDescriptor : declarativeShadowVariableDescriptors) {
                 var entityClass = declarativeShadowVariableDescriptor.getEntityDescriptor().getEntityClass();
                 if (entityClass.isInstance(entity)) {
                     var variableId = declarativeShadowVariableDescriptor.getVariableMetaModel();
-                    var updater = variableIdToUpdater.computeIfAbsent(variableId, ignored -> new VariableUpdaterInfo<>(
-                            variableId,
-                            declarativeShadowVariableDescriptor,
-                            declarativeShadowVariableDescriptor.getEntityDescriptor().getShadowVariableLoopedDescriptor(),
-                            declarativeShadowVariableDescriptor.getMemberAccessor(),
-                            declarativeShadowVariableDescriptor.getCalculator()::executeGetter));
-                    graph.addVariableReferenceEntity(entity, updater);
+                    var updaters = variableIdToUpdaters.computeUpdatersForVariableOnEntity(variableId,
+                            entity,
+                            () -> entityVariableToUpdatersMapper.apply(entity, declarativeShadowVariableDescriptor,
+                                    variableId));
+                    graph.addVariableReferenceEntity(entity, updaters);
                     for (var sourceRoot : declarativeShadowVariableDescriptor.getSources()) {
                         for (var source : sourceRoot.variableSourceReferences()) {
                             if (source.downstreamDeclarativeVariableMetamodel() != null) {

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

@@ -1,9 +1,12 @@
 package ai.timefold.solver.core.impl.domain.variable.declarative;
 
+import java.util.List;
+
 import org.jspecify.annotations.NullMarked;
 
 @NullMarked
-public record EntityVariablePair<Solution_>(Object entity, VariableUpdaterInfo<Solution_> variableReference, int graphNodeId) {
+public record EntityVariablePair<Solution_>(Object entity, List<VariableUpdaterInfo<Solution_>> variableReferences,
+        int graphNodeId) {
     @Override
     public boolean equals(Object object) {
         if (!(object instanceof EntityVariablePair<?> that))
@@ -18,6 +21,6 @@ public int hashCode() {
 
     @Override
     public String toString() {
-        return entity + ":" + variableReference.id();
+        return entity + ":" + variableReferences.stream().map(VariableUpdaterInfo::id).toList();
     }
 }