feat: add constraint stream profiling

benchmark/src/main/resources/benchmark.xsd

@@ -458,6 +458,9 @@
           <xs:element minOccurs="0" name="constraintStreamAutomaticNodeSharing" type="xs:boolean"/>
 
 
+          <xs:element minOccurs="0" name="constraintStreamProfilingEnabled" type="xs:boolean"/>
+
+
           <xs:element minOccurs="0" name="incrementalScoreCalculatorClass" type="xs:string"/>
 
 

core/src/build/revapi-differences.json

@@ -477,6 +477,17 @@
           "new": "class ai.timefold.solver.core.api.solver.event.BestSolutionChangedEvent<Solution_>",
           "interface": "java.io.Serializable",
           "justification": "BestSolutionChangedEvent no longer extends java.util.EventObject"
+        },
+        {
+          "ignore": true,
+          "code": "java.annotation.attributeValueChanged",
+          "old": "class ai.timefold.solver.core.config.score.director.ScoreDirectorFactoryConfig",
+          "new": "class ai.timefold.solver.core.config.score.director.ScoreDirectorFactoryConfig",
+          "annotationType": "jakarta.xml.bind.annotation.XmlType",
+          "attribute": "propOrder",
+          "oldValue": "{\"easyScoreCalculatorClass\", \"easyScoreCalculatorCustomProperties\", \"constraintProviderClass\", \"constraintProviderCustomProperties\", \"constraintStreamImplType\", \"incrementalScoreCalculatorClass\", \"incrementalScoreCalculatorCustomProperties\", \"scoreDrlList\", \"initializingScoreTrend\", \"assertionScoreDirectorFactory\"}",
+          "newValue": "{\"easyScoreCalculatorClass\", \"easyScoreCalculatorCustomProperties\", \"constraintProviderClass\", \"constraintProviderCustomProperties\", \"constraintStreamImplType\", \"constraintStreamAutomaticNodeSharing\", \"constraintStreamProfilingEnabled\", \"incrementalScoreCalculatorClass\", \"incrementalScoreCalculatorCustomProperties\", \"scoreDrlList\", \"initializingScoreTrend\", \"assertionScoreDirectorFactory\"}",
+          "justification": "Add constraint profiling to config"
         }
       ]
     }

core/src/main/java/ai/timefold/solver/core/config/score/director/ScoreDirectorFactoryConfig.java

@@ -27,6 +27,7 @@
         "constraintProviderCustomProperties",
         "constraintStreamImplType",
         "constraintStreamAutomaticNodeSharing",
+        "constraintStreamProfilingEnabled",
         "incrementalScoreCalculatorClass",
         "incrementalScoreCalculatorCustomProperties",
         "scoreDrlList",
@@ -46,6 +47,7 @@ public class ScoreDirectorFactoryConfig extends AbstractConfig<ScoreDirectorFact
     protected Map<String, String> constraintProviderCustomProperties = null;
     protected ConstraintStreamImplType constraintStreamImplType;
     protected Boolean constraintStreamAutomaticNodeSharing;
+    protected Boolean constraintStreamProfilingEnabled;
 
     protected Class<? extends IncrementalScoreCalculator> incrementalScoreCalculatorClass = null;
 
@@ -126,6 +128,14 @@ public void setConstraintStreamAutomaticNodeSharing(@Nullable Boolean constraint
         this.constraintStreamAutomaticNodeSharing = constraintStreamAutomaticNodeSharing;
     }
 
+    public Boolean getConstraintStreamProfilingEnabled() {
+        return constraintStreamProfilingEnabled;
+    }
+
+    public void setConstraintStreamProfilingEnabled(Boolean constraintStreamProfilingEnabled) {
+        this.constraintStreamProfilingEnabled = constraintStreamProfilingEnabled;
+    }
+
     public @Nullable Class<? extends IncrementalScoreCalculator> getIncrementalScoreCalculatorClass() {
         return incrementalScoreCalculatorClass;
     }
@@ -227,6 +237,12 @@ public void setAssertionScoreDirectorFactory(@Nullable ScoreDirectorFactoryConfi
         return this;
     }
 
+    public @NonNull ScoreDirectorFactoryConfig
+            withConstraintStreamProfilingEnabled(Boolean constraintStreamProfiling) {
+        this.constraintStreamProfilingEnabled = constraintStreamProfiling;
+        return this;
+    }
+
     public @NonNull ScoreDirectorFactoryConfig
             withIncrementalScoreCalculatorClass(
                     @NonNull Class<? extends IncrementalScoreCalculator> incrementalScoreCalculatorClass) {
@@ -288,6 +304,8 @@ public ScoreDirectorFactoryConfig withScoreDrls(String... scoreDrls) {
                 constraintStreamImplType, inheritedConfig.getConstraintStreamImplType());
         constraintStreamAutomaticNodeSharing = ConfigUtils.inheritOverwritableProperty(constraintStreamAutomaticNodeSharing,
                 inheritedConfig.getConstraintStreamAutomaticNodeSharing());
+        constraintStreamProfilingEnabled = ConfigUtils.inheritOverwritableProperty(constraintStreamProfilingEnabled,
+                inheritedConfig.getConstraintStreamProfilingEnabled());
         incrementalScoreCalculatorClass = ConfigUtils.inheritOverwritableProperty(
                 incrementalScoreCalculatorClass, inheritedConfig.getIncrementalScoreCalculatorClass());
         incrementalScoreCalculatorCustomProperties = ConfigUtils.inheritMergeableMapProperty(

core/src/main/java/ai/timefold/solver/core/impl/bavet/AbstractSession.java

@@ -63,4 +63,8 @@ public void settle() {
         nodeNetwork.settle();
     }
 
+    public final void summarizeProfileIfPresent() {
+        nodeNetwork.summarizeProfileIfPresent();
+    }
+
 }

core/src/main/java/ai/timefold/solver/core/impl/bavet/NodeNetwork.java

@@ -8,8 +8,12 @@
 
 import ai.timefold.solver.core.api.domain.solution.PlanningSolution;
 import ai.timefold.solver.core.impl.bavet.common.BavetRootNode;
+import ai.timefold.solver.core.impl.bavet.common.ConstraintProfiler;
 import ai.timefold.solver.core.impl.bavet.common.Propagator;
 
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
 /**
  * Represents Bavet's network of nodes, specific to a particular session.
  * Nodes only used by disabled constraints have already been removed.
@@ -19,10 +23,11 @@
  * @param layeredNodes nodes grouped first by their layer, then by their index within the layer;
  *        propagation needs to happen in this order.
  */
+@NullMarked
 public record NodeNetwork(Map<Class<?>, List<BavetRootNode<?>>> declaredClassToNodeMap,
-        Propagator[][] layeredNodes) {
+        Propagator[][] layeredNodes, @Nullable ConstraintProfiler constraintProfiler) {
 
-    public static final NodeNetwork EMPTY = new NodeNetwork(Map.of(), new Propagator[0][0]);
+    public static final NodeNetwork EMPTY = new NodeNetwork(Map.of(), new Propagator[0][0], null);
 
     public int forEachNodeCount() {
         return declaredClassToNodeMap.size();
@@ -70,6 +75,12 @@ private static void settleLayer(Propagator[] nodesInLayer) {
         }
     }
 
+    public void summarizeProfileIfPresent() {
+        if (constraintProfiler != null) {
+            constraintProfiler.summarize();
+        }
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o)

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

@@ -26,6 +26,11 @@ protected AbstractFlattenNode(int flattenStoreIndex, TupleLifecycle<OutTuple_> n
         this.propagationQueue = new StaticPropagationQueue<>(nextNodesTupleLifecycle);
     }
 
+    @Override
+    public final LifecycleKind lifecycleKind() {
+        return LifecycleKind.FLATTEN;
+    }
+
     @Override
     public final void insert(InTuple_ tuple) {
         if (tuple.getStore(flattenStoreIndex) != null) {

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

@@ -92,6 +92,11 @@ protected AbstractGroupNode(int groupStoreIndex,
                 environmentMode);
     }
 
+    @Override
+    public final LifecycleKind lifecycleKind() {
+        return LifecycleKind.GROUP_BY;
+    }
+
     @Override
     public final void insert(InTuple_ tuple) {
         if (tuple.getStore(groupStoreIndex) != null) {

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

@@ -18,6 +18,11 @@ protected AbstractMapNode(int inputStoreIndex, TupleLifecycle<OutTuple_> nextNod
         this.propagationQueue = new StaticPropagationQueue<>(nextNodesTupleLifecycle);
     }
 
+    @Override
+    public final LifecycleKind lifecycleKind() {
+        return LifecycleKind.MAP;
+    }
+
     @Override
     public final void insert(InTuple_ tuple) {
         if (tuple.getStore(inputStoreIndex) != null) {

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

@@ -1,14 +1,23 @@
 package ai.timefold.solver.core.impl.bavet.common;
 
+import java.util.Collections;
+import java.util.LinkedHashSet;
+import java.util.Set;
+
 import ai.timefold.solver.core.impl.score.stream.bavet.BavetConstraintSession;
 
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
 /**
  * @see PropagationQueue Description of the propagation mechanism.
  */
+@NullMarked
 public abstract class AbstractNode {
 
     private long id;
     private long layerIndex = -1;
+    private @Nullable Set<ConstraintNodeLocation> locationSet;
 
     /**
      * Instead of calling the propagation directly from here,
@@ -29,6 +38,20 @@ public final void setId(long id) {
         this.id = id;
     }
 
+    public Set<ConstraintNodeLocation> getLocationSet() {
+        if (locationSet == null) {
+            return Collections.emptySet();
+        }
+        return locationSet;
+    }
+
+    public void addLocationSet(Set<ConstraintNodeLocation> locationSet) {
+        if (this.locationSet == null) {
+            this.locationSet = new LinkedHashSet<>();
+        }
+        this.locationSet.addAll(locationSet);
+    }
+
     public final void setLayerIndex(long layerIndex) {
         if (layerIndex < 0) {
             throw new IllegalArgumentException("Impossible state: layer index (" + layerIndex + ") must be at least 0.");

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

@@ -1,8 +1,10 @@
 package ai.timefold.solver.core.impl.bavet.common;
 
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -14,26 +16,39 @@
 import ai.timefold.solver.core.impl.bavet.NodeNetwork;
 import ai.timefold.solver.core.impl.bavet.common.tuple.InOutTupleStorePositionTracker;
 import ai.timefold.solver.core.impl.bavet.common.tuple.LeftTupleLifecycle;
+import ai.timefold.solver.core.impl.bavet.common.tuple.ProfilingTupleLifecycle;
 import ai.timefold.solver.core.impl.bavet.common.tuple.RightTupleLifecycle;
 import ai.timefold.solver.core.impl.bavet.common.tuple.Tuple;
 import ai.timefold.solver.core.impl.bavet.common.tuple.TupleLifecycle;
+import ai.timefold.solver.core.impl.score.stream.bavet.common.Scorer;
 
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
+@NullMarked
 public abstract class AbstractNodeBuildHelper<Stream_ extends BavetStream> {
 
     private final Set<Stream_> activeStreamSet;
     private final Map<AbstractNode, Stream_> nodeCreatorMap;
     private final Map<Stream_, TupleLifecycle<? extends Tuple>> tupleLifecycleMap;
     private final Map<Stream_, Integer> storeIndexMap;
 
+    @Nullable
+    private final ConstraintProfiler constraintProfiler;
+
+    @Nullable
     private List<AbstractNode> reversedNodeList;
+    private long nextLifecycleProfilingId = 0;
 
-    protected AbstractNodeBuildHelper(Set<Stream_> activeStreamSet) {
+    protected AbstractNodeBuildHelper(Set<Stream_> activeStreamSet,
+            @Nullable ConstraintProfiler constraintProfiler) {
         this.activeStreamSet = activeStreamSet;
         int activeStreamSetSize = activeStreamSet.size();
         this.nodeCreatorMap = new HashMap<>(Math.max(16, activeStreamSetSize));
         this.tupleLifecycleMap = new HashMap<>(Math.max(16, activeStreamSetSize));
         this.storeIndexMap = new HashMap<>(Math.max(16, activeStreamSetSize / 2));
         this.reversedNodeList = new ArrayList<>(activeStreamSetSize);
+        this.constraintProfiler = constraintProfiler;
     }
 
     public boolean isStreamActive(Stream_ stream) {
@@ -46,6 +61,7 @@ public void addNode(AbstractNode node, Stream_ creator) {
 
     public void addNode(AbstractNode node, Stream_ creator, Stream_ parent) {
         reversedNodeList.add(node);
+        node.addLocationSet(creator.getLocationSet());
         nodeCreatorMap.put(node, creator);
         if (!(node instanceof BavetRootNode<?>)) {
             if (parent == null) {
@@ -58,13 +74,40 @@ public void addNode(AbstractNode node, Stream_ creator, Stream_ parent) {
 
     public void addNode(AbstractNode node, Stream_ creator, Stream_ leftParent, Stream_ rightParent) {
         reversedNodeList.add(node);
+        node.addLocationSet(creator.getLocationSet());
         nodeCreatorMap.put(node, creator);
         putInsertUpdateRetract(leftParent, TupleLifecycle.ofLeft((LeftTupleLifecycle<? extends Tuple>) node));
         putInsertUpdateRetract(rightParent, TupleLifecycle.ofRight((RightTupleLifecycle<? extends Tuple>) node));
     }
 
     public <Tuple_ extends Tuple> void putInsertUpdateRetract(Stream_ stream,
             TupleLifecycle<Tuple_> tupleLifecycle) {
+        if (constraintProfiler != null) {
+            tupleLifecycle = TupleLifecycle.profiling(constraintProfiler, nextLifecycleProfilingId,
+                    stream, tupleLifecycle);
+
+            // Profiling wraps each node inside a profiler node, so we need to unwrap the profiling node first
+            if (tupleLifecycle instanceof ProfilingTupleLifecycle<Tuple_> profilingTupleLifecycle &&
+                    profilingTupleLifecycle.delegate() instanceof Scorer<Tuple_> scorer) {
+                // This is a score, so we can navigate up its parents
+                // to find all locations corresponding to this constraint
+                var queue = new ArrayDeque<BavetStream>();
+                var constraintLocationSet = new LinkedHashSet<ConstraintNodeLocation>();
+                queue.add(stream);
+                while (!queue.isEmpty()) {
+                    var currentStream = queue.poll();
+                    constraintLocationSet.addAll(currentStream.getLocationSet());
+                    if (currentStream instanceof BavetStreamBinaryOperation<?> binaryOperation) {
+                        queue.add(binaryOperation.getLeftParent());
+                        queue.add(binaryOperation.getRightParent());
+                    } else if (currentStream.getParent() != null) {
+                        queue.add(currentStream.getParent());
+                    }
+                }
+                constraintProfiler.registerConstraint(scorer.getConstraintRef(), constraintLocationSet);
+            }
+            nextLifecycleProfilingId++;
+        }
         tupleLifecycleMap.put(stream, tupleLifecycle);
     }
 
@@ -153,19 +196,30 @@ public AbstractNode findParentNode(Stream_ childNodeCreator) {
     }
 
     public static NodeNetwork buildNodeNetwork(List<AbstractNode> nodeList,
-            Map<Class<?>, List<BavetRootNode<?>>> declaredClassToNodeMap) {
+            Map<Class<?>, List<BavetRootNode<?>>> declaredClassToNodeMap,
+            AbstractNodeBuildHelper<?> nodeBuildHelper) {
         var layerMap = new TreeMap<Long, List<Propagator>>();
+        var profiler = nodeBuildHelper.constraintProfiler;
         for (var node : nodeList) {
-            layerMap.computeIfAbsent(node.getLayerIndex(), k -> new ArrayList<>())
-                    .add(node.getPropagator());
+            var layer = node.getLayerIndex();
+            var propagator = node.getPropagator();
+            if (profiler != null) {
+                var profileKey = nodeBuildHelper.nextLifecycleProfilingId;
+                nodeBuildHelper.nextLifecycleProfilingId++;
+                var profileId = new ConstraintNodeProfileId(profileKey, LifecycleKind.PROPAGATION, node.getLocationSet());
+                nodeBuildHelper.constraintProfiler.register(profileId);
+                propagator = new ProfilingPropagator(profiler, profileId, propagator);
+            }
+            layerMap.computeIfAbsent(layer, k -> new ArrayList<>())
+                    .add(propagator);
         }
         var layerCount = layerMap.size();
         var layeredNodes = new Propagator[layerCount][];
         for (var i = 0; i < layerCount; i++) {
             var layer = layerMap.get((long) i);
             layeredNodes[i] = layer.toArray(new Propagator[0]);
         }
-        return new NodeNetwork(declaredClassToNodeMap, layeredNodes);
+        return new NodeNetwork(declaredClassToNodeMap, layeredNodes, nodeBuildHelper.constraintProfiler);
     }
 
     public <BuildHelper_ extends AbstractNodeBuildHelper<Stream_>> List<AbstractNode> buildNodeList(Set<Stream_> streamSet,