Add keep argument for Graph::remove_unused_nodes()

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

@@ -189,6 +189,13 @@ class Graph {
     ///
     /// Returns the number of nodes removed from the graph.
     ssize_t remove_unused_nodes(bool ignore_listeners = false);
+    ssize_t remove_unused_nodes(std::span<Node*> keep, bool ignore_listeners = false);
+    ssize_t remove_unused_nodes(
+        std::ranges::contiguous_range auto&& keep,
+        bool ignore_listeners = false
+    ) {
+        return remove_unused_nodes(std::span<Node*>(keep), ignore_listeners);
+    }
 
     /// Call revert on every `Node` in the `Graph`.
     void revert(State& state) const;
@@ -355,7 +362,7 @@ class Node {
     template <class NodeType, class... Args>
     friend NodeType* Graph::emplace_node(Args&&...);
     friend ssize_t Graph::remove_redundant_nodes(bool, double);
-    friend ssize_t Graph::remove_unused_nodes(bool);
+    friend ssize_t Graph::remove_unused_nodes(std::span<Node*>, bool);
     friend void Graph::reset_topological_sort();
     friend void Graph::swap_decisions(DecisionNode* x_ptr, DecisionNode* y_ptr);
     friend void Graph::topological_sort();

dwave/optimization/src/graph.cpp

@@ -303,20 +303,20 @@ ssize_t Graph::remove_redundant_nodes(bool ignore_listeners, double time_limit_s
     // already done.
     // They are arbitrary values, but we steer away from -1 because that's used to
     // indicate unsorted
-    constexpr ssize_t seen = -2;  // already checked for duplicates
-    constexpr ssize_t drop = -3;  // marked for deletion
+    constexpr ssize_t SEEN = -2;  // already checked for duplicates
+    constexpr ssize_t DROP = -3;  // marked for deletion
 
     // The actions that we always want to do before returning, whether it's because
     // we hit the time limit or because we have no more work to do.
     auto cleanup = [&]() -> ssize_t {
         // we should only ever have swapped the objective_ptr
-        assert(objective_ptr_ == nullptr or objective_ptr_->topological_index() != drop);
+        assert(objective_ptr_ == nullptr or objective_ptr_->topological_index() != DROP);
 
         // Whether a node was dropped
-        auto dropped = [&drop, &ignore_listeners](const auto& ptr) {
+        auto dropped = [&DROP, &ignore_listeners](const auto& ptr) {
             if (not ignore_listeners and ptr->num_listeners() > 0) return false;
-            assert(ptr->topological_index_ != drop or ptr->successors_.empty());
-            return ptr->topological_index() == drop;
+            assert(ptr->topological_index_ != DROP or ptr->successors_.empty());
+            return ptr->topological_index() == DROP;
         };
 
         // Go through our "special" nodes and drop anything
@@ -325,7 +325,7 @@ ssize_t Graph::remove_redundant_nodes(bool ignore_listeners, double time_limit_s
         std::erase_if(constants_, dropped);
 
         std::erase_if(constraints_, dropped);
-        assert(objective_ptr_ == nullptr or objective_ptr_->topological_index_ != drop);
+        assert(objective_ptr_ == nullptr or objective_ptr_->topological_index_ != DROP);
 
         // This is the step that actually deallocates the node
         for (auto& uptr : nodes_) {
@@ -363,7 +363,7 @@ ssize_t Graph::remove_redundant_nodes(bool ignore_listeners, double time_limit_s
         to_ptr->take_successors(*from_ptr);
 
         // Mark from for dropping
-        from_ptr->topological_index_ = drop;
+        from_ptr->topological_index_ = DROP;
 
         // We also want to fix the objective if relevant
         if (objective_ptr_ != nullptr and static_cast<Node*>(objective_ptr_) == from_ptr) {
@@ -382,10 +382,10 @@ ssize_t Graph::remove_redundant_nodes(bool ignore_listeners, double time_limit_s
     // The first set of nodes we're worried about are the constants - the only
     // class of root node that can have redundancy
     for (ssize_t i = 0, num_constants = constants_.size(); i < num_constants; ++i) {
-        if (constants_[i]->topological_index_ == drop) continue;
+        if (constants_[i]->topological_index_ == DROP) continue;
 
         for (ssize_t j = i + 1; j < num_constants; ++j) {
-            if (constants_[j]->topological_index_ == drop) continue;
+            if (constants_[j]->topological_index_ == DROP) continue;
 
             // the topological indices of our constants should not yet be touched
             // and because they are added in order, they should always be ascending
@@ -403,7 +403,7 @@ ssize_t Graph::remove_redundant_nodes(bool ignore_listeners, double time_limit_s
             transfer(constants_[j], constants_[i]);
         }
 
-        constants_[i]->topological_index_ = seen;
+        constants_[i]->topological_index_ = SEEN;
     }
 
     // For each node, we check pairwise among all of its successors.
@@ -443,23 +443,28 @@ ssize_t Graph::remove_redundant_nodes(bool ignore_listeners, double time_limit_s
 
             // Ok, we've checked i against everything, so if we didn't drop it
             // we can mark it as seen
-            if (successors[i]->topological_index_ >= 0) successors[i]->topological_index_ = seen;
+            if (successors[i]->topological_index_ >= 0) successors[i]->topological_index_ = SEEN;
         }
     }
 
     return cleanup();
 }
 
+
 ssize_t Graph::remove_unused_nodes(bool ignore_listeners) {
+    return remove_unused_nodes({}, ignore_listeners);
+}
+
+ssize_t Graph::remove_unused_nodes(std::span<Node*> keep, bool ignore_listeners) {
     if (topologically_sorted_) throw std::logic_error("cannot remove nodes from a locked model");
 
     // Establish a topological ordering. We'll use the fact that the node list
     // is ordered, but we're going to mess with the topological indices!
     topological_sort();
 
     // Specifically we'll be marking the nodes with either
-    constexpr ssize_t keep = -2;  // always keep
-    constexpr ssize_t drop = -3;  // to be dropped
+    constexpr ssize_t KEEP = -2;  // always keep
+    constexpr ssize_t DROP = -3;  // to be dropped
 
     // We'll use this later to calculate how many nodes we removed
     const ssize_t num_nodes = this->num_nodes();
@@ -469,30 +474,33 @@ ssize_t Graph::remove_unused_nodes(bool ignore_listeners) {
     // nodes.
 
     // First up our roots - specifically the decisions are always kept.
-    for (auto* ptr : decisions_) ptr->topological_index_ = keep;
+    for (auto* ptr : decisions_) ptr->topological_index_ = KEEP;
 
     // Constants and Inputs are allowed to be removed
 
     // Next up our important leaves.
-    for (auto* ptr : constraints_) ptr->topological_index_ = keep;
-    if (objective_ptr_) objective_ptr_->topological_index_ = keep;
+    for (auto* ptr : constraints_) ptr->topological_index_ = KEEP;
+    if (objective_ptr_) objective_ptr_->topological_index_ = KEEP;
 
     // Some nodes are just not allowed to be removed for other reasons
     for (auto& uptr : nodes_) {
-        if (not uptr->removable_()) uptr->topological_index_ = keep;
+        if (not uptr->removable_()) uptr->topological_index_ = KEEP;
     }
 
+    // And of course, any node that the user has explicitly told us we must keep
+    for (auto* ptr : keep) ptr->topological_index_ = KEEP;
+
     // Also, any nodes that have a listener, that is a node is holding a Node::expired_ptr()
     if (not ignore_listeners) {
         for (auto& uptr : nodes_) {
-            if (uptr->num_listeners() > 0) uptr->topological_index_ = keep;
+            if (uptr->num_listeners() > 0) uptr->topological_index_ = KEEP;
         }
     }
 
     // Having established what we're keeping, time to mark stuff for removal!
 
     for (auto& uptr : nodes_ | std::views::reverse) {
-        if (uptr->topological_index_ == keep) continue;  // we marked these to keep
+        if (uptr->topological_index_ == KEEP) continue;  // we marked these to keep
         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.
@@ -505,17 +513,17 @@ ssize_t Graph::remove_unused_nodes(bool ignore_listeners) {
             assert(num_removed > 0);
         }
 
-        uptr->topological_index_ = drop;
+        uptr->topological_index_ = DROP;
     }
 
     // Now let's start actually dropping stuff
-    auto dropped = [&drop](const auto& ptr) { return ptr->topological_index_ == drop; };
+    auto dropped = [&DROP](const auto& ptr) { return ptr->topological_index_ == DROP; };
 
     assert(std::ranges::none_of(decisions_, dropped));
     std::erase_if(inputs_, dropped);
     std::erase_if(constants_, dropped);
 
-    assert(objective_ptr_ == nullptr or objective_ptr_->topological_index_ != drop);
+    assert(objective_ptr_ == nullptr or objective_ptr_->topological_index_ != DROP);
     assert(std::ranges::none_of(constraints_, dropped));
 
     // Traverse the nodes_ one last time, this is what actualy deallocates the nodes.

releasenotes/notes/remove-unused-nodes-keeps-857275672e5db761.yaml

@@ -0,0 +1,3 @@
+---
+features:
+  - Add ``keep`` argument for ``Graph::remove_unused_nodes()``.

tests/cpp/test_graph.cpp

@@ -493,7 +493,8 @@ TEST_CASE("Graph::remove_unused_nodes()") {
         }
 
         WHEN("We add two successors that are not used in a constraint") {
-            graph.emplace_node<LogicalNode>(graph.emplace_node<AbsoluteNode>(i_ptr));
+            auto* abs_ptr = graph.emplace_node<AbsoluteNode>(i_ptr);
+            auto* logical_ptr = graph.emplace_node<LogicalNode>(abs_ptr);
 
             THEN("remove_unused_nodes() removes them") {
                 ssize_t num_removed = graph.remove_unused_nodes();
@@ -502,6 +503,38 @@ TEST_CASE("Graph::remove_unused_nodes()") {
                 CHECK(i_ptr->topological_index() == 0);
                 CHECK(i_ptr->successors().size() == 0);
             }
+
+            THEN("remove_unused_nodes({logical_ptr}) keeps both") {
+                auto keep = std::vector<Node*>{logical_ptr};
+
+                WHEN("keep is given as an rvalue") {
+                    ssize_t num_removed = graph.remove_unused_nodes(std::move(keep));
+                    CHECK(num_removed == 0);
+                    CHECK(graph.num_nodes() == 3);
+                }
+
+                WHEN("keep is given as an lvalue") {
+                    ssize_t num_removed = graph.remove_unused_nodes(keep);
+                    CHECK(num_removed == 0);
+                    CHECK(graph.num_nodes() == 3);
+                }
+            }
+
+            THEN("remove_unused_nodes({abs_ptr}) keeps one") {
+                auto keep = std::vector<Node*>{abs_ptr};
+
+                WHEN("keep is given as an rvalue") {
+                    ssize_t num_removed = graph.remove_unused_nodes(std::move(keep));
+                    CHECK(num_removed == 1);
+                    CHECK(graph.num_nodes() == 2);
+                }
+
+                WHEN("keep is given as an lvalue") {
+                    ssize_t num_removed = graph.remove_unused_nodes(keep);
+                    CHECK(num_removed == 1);
+                    CHECK(graph.num_nodes() == 2);
+                }
+            }
         }
 
         WHEN("We add two successors and use them in a constraint") {