[WIP] Add Graph::mutated() (extra NodeStateData implementation)

dwave/optimization/include/dwave-optimization/graph.hpp

@@ -115,6 +115,25 @@ class Graph {
     std::span<InputNode* const> inputs() noexcept { return inputs_; }
     std::span<const InputNode* const> inputs() const noexcept { return inputs_; }
 
+    /// Return the decision nodes that have been "mutated", meaning that they
+    /// have pending changes that must be propagated before committing or
+    /// reverting. Equivalently, the combined descendants of the returned nodes
+    /// are guaranteed to be a superset of the nodes that require having
+    /// `propagate()` and then `commit()` or `revert()` called on them.
+    ///
+    /// Notes:
+    ///  - Using this method in conjunction with `descendants()` and
+    ///  `propagate()`/`commit()`/`revert()` may be inefficient compared to
+    ///  doing these manually in the case where you know only some descendants
+    ///  of one or more of the decision nodes are relevant, e.g. only one of the
+    ///  `DisjointListNode` successors of `DisjointListsNode` has pending
+    ///  changes and the rest of the `DisjointListNode`s (and their descendants)
+    ///  can be ignored
+    ///  - This method will return the same nodes before and after calling
+    ///  `propagate()`. Only after committing/reverting will the returned list
+    ///  be empty again.
+    std::span<const DecisionNode*> mutated(State& state) const;
+
     /// All of the nodes in the graph.
     std::span<const std::unique_ptr<Node>> nodes() const { return nodes_; }
 
@@ -139,7 +158,7 @@ class Graph {
     void propagate(State& state) const;
 
     /// Call the propagate method on each node in changed. Note this does not call propagate on
-    /// the descendents of changed.
+    /// the descendants of changed.
     void propagate(State& state, std::span<const Node*> changed) const;
     void propagate(State& state, std::vector<const Node*>&& changed) const;
 

dwave/optimization/src/graph.cpp

@@ -27,6 +27,7 @@
 #endif
 
 #include "dwave-optimization/array.hpp"
+#include "dwave-optimization/nodes/collections.hpp"
 #include "dwave-optimization/nodes/constants.hpp"
 #include "dwave-optimization/nodes/inputs.hpp"
 
@@ -110,7 +111,7 @@ std::vector<const Node*> Graph::descendants(std::vector<const Node*> sources) co
     return descendants(state, sources);
 }
 
-State Graph::empty_state() const { return State(num_nodes()); }
+State Graph::empty_state() const { return State(num_nodes() + 1); }
 
 State Graph::empty_state() {
     topological_sort();
@@ -129,6 +130,16 @@ bool Graph::feasible(const State& state) const {
     return true;
 }
 
+class GraphStateData : public NodeStateData {
+ public:
+    virtual std::unique_ptr<NodeStateData> copy() const override {
+        assert(typeid(*this) == typeid(NodeStateData) && "subclasses should overload copy()");
+        return std::make_unique<GraphStateData>(*this);
+    }
+
+    std::vector<const DecisionNode*> mutated_nodes;
+};
+
 State Graph::initialize_state() const {
     auto state = empty_state();
     initialize_state(state);
@@ -141,21 +152,55 @@ State Graph::initialize_state() {
 }
 
 void Graph::initialize_state(State& state) const {
-    assert(static_cast<int>(state.size()) == num_nodes() and "unexpected state length");
+    assert(static_cast<ssize_t>(state.size()) == num_nodes() + 1 and "unexpected state length");
     assert(topologically_sorted_ and "graph must be topologically sorted");
 
     for (int i = 0, end = num_nodes(); i < end; ++i) {
         if (state[i]) continue;  // should this clear any pending changes?
 
         nodes_[i]->initialize_state(state);
     }
+
+    state.back() = std::make_unique<GraphStateData>();
 }
 
 void Graph::initialize_state(State& state) {
     topological_sort();
     static_cast<const Graph*>(this)->initialize_state(state);
 }
 
+std::span<const DecisionNode*> Graph::mutated(State& state) const {
+    assert(static_cast<ssize_t>(state.size()) == num_nodes() + 1);
+
+    auto& mutated_nodes = static_cast<GraphStateData*>(state.back().get())->mutated_nodes;
+
+    mutated_nodes.clear();
+    for (const DecisionNode* dec_ptr : decisions()) {
+        if (const ArrayNode* arr_ptr = dynamic_cast<const ArrayNode*>(dec_ptr); arr_ptr) {
+            if (not arr_ptr->diff(state).empty()) {
+                mutated_nodes.push_back(dec_ptr);
+            }
+        } else if (
+            dynamic_cast<const DisjointListsNode*>(dec_ptr) or
+            dynamic_cast<const DisjointBitSetsNode*>(dec_ptr)
+        ) {
+            for (const Node* suc_ptr : dec_ptr->successors()) {
+                const ArrayNode* arr_ptr = dynamic_cast<const ArrayNode*>(suc_ptr);
+                assert(arr_ptr and "all successors should be array nodes");
+                if (not arr_ptr->diff(state).empty()) {
+                    mutated_nodes.push_back(dec_ptr);
+                    break;
+                }
+            }
+        } else {
+            assert(false and "unknown decision node type");
+            unreachable();
+        }
+    }
+
+    return mutated_nodes;
+}
+
 void Graph::propagate(State& state) const {
     std::ranges::for_each(nodes(), [&state](const auto& ptr) { ptr->propagate(state); });
 }
@@ -284,7 +329,7 @@ ssize_t Graph::remove_unused_nodes(bool ignore_listeners) {
 
     for (auto& uptr : nodes_ | std::views::reverse) {
         if (uptr->topological_index_ == keep) continue;  // we marked these to keep
-        if (uptr->successors().size() > 0) continue;  // this node is used by other nodes
+        if (uptr->successors().size() > 0) continue;     // this node is used by other nodes
 
         // We have a node with no successors and that we haven't marked it as important.
         // So let's mark it to be dropped later.

tests/cpp/test_graph.cpp

@@ -233,7 +233,9 @@ TEST_CASE("Graph constructors, assignment operators, and swapping") {
     }
 }
 
-TEST_CASE("Graph::commit(), Graph::descendants(), Graph::propagate(), and Graph::revert") {
+TEST_CASE(
+    "Graph::commit(), Graph::descendants(), Graph::mutated(), Graph::propagate(), and Graph::revert"
+) {
     auto graph = Graph();
     auto* x_ptr = graph.emplace_node<BinaryNode>();
     auto* y_ptr = graph.emplace_node<BinaryNode>();
@@ -246,17 +248,23 @@ TEST_CASE("Graph::commit(), Graph::descendants(), Graph::propagate(), and Graph:
         CHECK_THAT(descendants, Catch::Matchers::RangeEquals(std::vector<Node*>{x_ptr, z_ptr}));
     }
     SECTION("Propagate all") {
+        CHECK_THAT(graph.mutated(state), Catch::Matchers::RangeEquals(std::vector<Node*>{}));
+
         CHECK(x_ptr->view(state).front() == 0);
         CHECK(y_ptr->view(state).front() == 0);
         CHECK(z_ptr->view(state).front() == 0);
 
         x_ptr->flip(state, 0);
+        CHECK_THAT(graph.mutated(state), Catch::Matchers::RangeEquals({x_ptr}));
+
         y_ptr->flip(state, 0);
 
         CHECK(x_ptr->diff(state).size());
         CHECK(y_ptr->diff(state).size());
         CHECK(z_ptr->diff(state).empty());  // not yet propagated to
 
+        CHECK_THAT(graph.mutated(state), Catch::Matchers::RangeEquals({x_ptr, y_ptr}));
+
         graph.propagate(state);
 
         CHECK(x_ptr->view(state).front() == 1);
@@ -267,6 +275,8 @@ TEST_CASE("Graph::commit(), Graph::descendants(), Graph::propagate(), and Graph:
         CHECK(y_ptr->diff(state).size());
         CHECK(z_ptr->diff(state).size());  // now has pending changes
 
+        CHECK_THAT(graph.mutated(state), Catch::Matchers::RangeEquals({x_ptr, y_ptr}));
+
         SECTION("Commit all") {
             graph.commit(state);
 
@@ -278,6 +288,9 @@ TEST_CASE("Graph::commit(), Graph::descendants(), Graph::propagate(), and Graph:
             CHECK(x_ptr->diff(state).empty());
             CHECK(y_ptr->diff(state).empty());
             CHECK(z_ptr->diff(state).empty());
+
+            // Committing should reset the mutated nodes
+            CHECK_THAT(graph.mutated(state), Catch::Matchers::RangeEquals(std::vector<Node*>{}));
         }
 
         SECTION("Revert all") {
@@ -291,6 +304,9 @@ TEST_CASE("Graph::commit(), Graph::descendants(), Graph::propagate(), and Graph:
             CHECK(x_ptr->diff(state).empty());
             CHECK(y_ptr->diff(state).empty());
             CHECK(z_ptr->diff(state).empty());
+
+            // Reverting should reset the mutated nodes
+            CHECK_THAT(graph.mutated(state), Catch::Matchers::RangeEquals(std::vector<Node*>{}));
         }
     }
 }