Implement state checkpointing for decision nodes

.circleci/config.yml

@@ -30,7 +30,7 @@ environment: &global-environment
 jobs:
   python-linux:
     docker:
-      - image: cimg/python:3.13  # just need a version that can install cibuildwheel
+      - image: cimg/python:3.13  # need a version that can install cibuildwheel and that has docker
 
     environment:
       <<: *global-environment
@@ -57,7 +57,7 @@ jobs:
 
   python-linux-debug:
     docker:
-      - image: cimg/python:3.10
+      - image: python:3.10
 
     steps:
       - checkout
@@ -127,7 +127,7 @@ jobs:
 
   python-sdist:
     docker:
-      - image: cimg/python:3.10
+      - image: python:3.10
 
     steps:
       - checkout

dwave/optimization/include/dwave-optimization/array.hpp

@@ -347,6 +347,9 @@ struct Update {
     // Return true if the update does nothing - that is old and value are the same.
     bool identity() const { return null() || old == value; }
 
+    // Return the update that would undo the current update
+    Update inverse() const { return Update(index, value, old); }
+
     // Use NaN to represent the "nothing" value used in placements/removals
     static constexpr double nothing = std::numeric_limits<double>::signaling_NaN();
 

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

@@ -155,6 +155,11 @@ class Graph {
         std::function<bool(const Graph&, State&)> accept = [](const Graph&, State&) { return true; }
     ) const;
 
+    /// Propagate any pending changes to all nodes in the graph and commit them.
+    void propose(State& state) const;
+    // dev note: the name is a bit funny in this case, but we essentially want
+    // an overload for a "default" `sources` and `accept`.
+
     /// Initialize the state of the given node and all predecessors recursively.
     static void recursive_initialize(State& state, const Node* ptr);
     /// Reset the state of the given node and all successors recursively.
@@ -463,6 +468,13 @@ NodeType* Graph::emplace_node(Args&&... args) {
 class ArrayNode : public Array, public virtual Node {};
 class DecisionNode : public Decision, public virtual Node {
  public:
+    /// Set the current state to match the one at the time the given checkpoint was created.
+    virtual void assign_from_checkpoint(State& state, checkpoint_type& checkpoint) const = 0;
+    virtual void assign_from_checkpoint(State& state, checkpoint_type&& checkpoint) const = 0;
+
+    /// Get a checkpoint, an IOU that can be used to return the node to its current state.
+    virtual checkpoint_type checkpoint(State& state) const = 0;
+
     /// Decision nodes by definition do not have a deterministic state.
     bool deterministic_state() const final { return false; }
 

dwave/optimization/include/dwave-optimization/nodes/collections.hpp

@@ -32,8 +32,15 @@ class CollectionNode : public ArrayOutputMixin<ArrayNode>, public DecisionNode {
     // Set the node's state, tracking the diff.
     void assign(State& state, std::vector<double> values) const;
 
+    /// @copydoc DecisionNode::assign_from_checkpoint()
+    void assign_from_checkpoint(State& state, checkpoint_type& checkpoint) const override;
+    void assign_from_checkpoint(State& state, checkpoint_type&& checkpoint) const override;
+
     const double* buff(const State& state) const override;
 
+    /// @copydoc DecisionNode::checkpoint()
+    checkpoint_type checkpoint(State& state) const override;
+
     void commit(State&) const override;
 
     std::span<const Update> diff(const State& state) const override;
@@ -100,6 +107,13 @@ class DisjointBitSetsNode : public DecisionNode {
     // i.e. the set `range(primary_set_size)`.
     DisjointBitSetsNode(ssize_t primary_set_size, ssize_t num_disjoint_sets);
 
+    /// @copydoc DecisionNode::assign_from_checkpoint()
+    void assign_from_checkpoint(State& state, checkpoint_type& checkpoint) const override;
+    void assign_from_checkpoint(State& state, checkpoint_type&& checkpoint) const override;
+
+    /// @copydoc DecisionNode::checkpoint()
+    checkpoint_type checkpoint(State& state) const override;
+
     void commit(State&) const override;
 
     ssize_t get_containing_set_index(State& state, ssize_t element_i) const;
@@ -175,6 +189,13 @@ class DisjointListsNode : public DecisionNode {
     // i.e. the set `range(primary_set_size)`.
     DisjointListsNode(ssize_t primary_set_size, ssize_t num_disjoint_lists);
 
+    /// @copydoc DecisionNode::assign_from_checkpoint()
+    void assign_from_checkpoint(State& state, checkpoint_type& checkpoint) const override;
+    void assign_from_checkpoint(State& state, checkpoint_type&& checkpoint) const override;
+
+    /// @copydoc DecisionNode::checkpoint()
+    checkpoint_type checkpoint(State& state) const override;
+
     void commit(State&) const override;
 
     ssize_t get_disjoint_list_size(State& state, ssize_t list_index) const;

dwave/optimization/include/dwave-optimization/nodes/numbers.hpp

@@ -17,7 +17,6 @@
 #include <optional>
 #include <ranges>
 #include <span>
-#include <utility>
 #include <vector>
 
 #include "dwave-optimization/array.hpp"
@@ -122,6 +121,9 @@ class NumberNode : public ArrayOutputMixin<ArrayNode>, public DecisionNode {
 
     // NumberNode methods *****************************************************
 
+    /// @copydoc DecisionNode::checkpoint()
+    checkpoint_type checkpoint(State& state) const override;
+
     // In the given state, swap the value of index i with the value of index j.
     // Users may pass the slices (per sum constraint) that each index lies on.
     void exchange(
@@ -290,6 +292,10 @@ class IntegerNode : public NumberNode {
 
     // IntegerNode methods ****************************************************
 
+    /// @copydoc DecisionNode::assign_from_checkpoint()
+    void assign_from_checkpoint(State& state, checkpoint_type& checkpoint) const override;
+    void assign_from_checkpoint(State& state, checkpoint_type&& checkpoint) const override;
+
     // Set the value at the given index in the given state.
     // Users may pass the slices (per sum constraint) that each index lies on.
     void set_value(

dwave/optimization/include/dwave-optimization/nodes/testing.hpp

@@ -105,6 +105,28 @@ class DynamicArrayTestingNode : public ArrayOutputMixin<ArrayNode>, public Decis
     void revert(State&) const override;
     void update(State&, int) const override;
 
+    // Overloads required by the DecisionNode ABC *****************************
+
+    // DynamicArrayTestingNode does not implement checkpointing
+    [[noreturn]] void assign_from_checkpoint(
+        State& state,
+        checkpoint_type& checkpoint
+    ) const override {
+        assert(false and "not implemented");
+        unreachable();
+    }
+    [[noreturn]] void assign_from_checkpoint(
+        State& state,
+        checkpoint_type&& checkpoint
+    ) const override {
+        assert(false and "not implemented");
+        unreachable();
+    }
+    [[noreturn]] virtual checkpoint_type checkpoint(State& state) const override {
+        assert(false and "not implemented");
+        unreachable();
+    }
+
     // State mutation methods *************************************************
 
     // Grow the array by a single row of the given values.

dwave/optimization/include/dwave-optimization/state.hpp

@@ -36,4 +36,17 @@ struct NodeStateData {
 
 using State = typename std::vector<std::unique_ptr<NodeStateData>>;
 
+/// A generic base class for node checkpoints.
+struct NodeStateCheckpoint {
+    NodeStateCheckpoint() = default;
+    NodeStateCheckpoint(const NodeStateCheckpoint&) = delete;
+    NodeStateCheckpoint(NodeStateCheckpoint&&) = delete;
+    NodeStateCheckpoint& operator=(const NodeStateCheckpoint&) = delete;
+    NodeStateCheckpoint& operator=(NodeStateCheckpoint&&) = delete;
+
+    virtual ~NodeStateCheckpoint() = default;
+};
+
+using checkpoint_type = std::unique_ptr<NodeStateCheckpoint>;
+
 }  // namespace dwave::optimization

dwave/optimization/src/graph.cpp

@@ -230,6 +230,11 @@ void Graph::propose(
     }
 }
 
+void Graph::propose(State& state) const {
+    propagate(state);
+    commit(state);
+}
+
 void Graph::recursive_initialize(State& state, const Node* ptr) {
     ssize_t index = ptr->topological_index();
 

dwave/optimization/src/nodes/_checkpoints.cpp

@@ -0,0 +1,91 @@
+// Copyright 2026 D-Wave
+//
+//    Licensed under the Apache License, Version 2.0 (the "License");
+//    you may not use this file except in compliance with the License.
+//    You may obtain a copy of the License at
+//
+//        http://www.apache.org/licenses/LICENSE-2.0
+//
+//    Unless required by applicable law or agreed to in writing, software
+//    distributed under the License is distributed on an "AS IS" BASIS,
+//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//    See the License for the specific language governing permissions and
+//    limitations under the License.
+
+#include "_checkpoints.hpp"
+
+namespace dwave::optimization {
+
+CheckpointableState::~CheckpointableState() {
+    if (prev_ptr_ == nullptr) return;  // nothing to clean up
+    prev_ptr_->next_ptr_ = static_cast<CheckpointableState*>(nullptr);
+}
+
+// Place self between the state and any checkpoint it's currently holding
+LinkedListCheckpoint::LinkedListCheckpoint(CheckpointableState& state) :
+    prev_ptr_(state.prev_ptr_), next_ptr_(&state) {
+    if (prev_ptr_ != nullptr) prev_ptr_->next_ptr_ = this;
+    state.prev_ptr_ = this;
+}
+
+LinkedListCheckpoint::~LinkedListCheckpoint() {
+    if (prev_ptr_ != nullptr) prev_ptr_->next_ptr_ = next_ptr_;
+
+    // Now make sure next_ptr is pointing to prev_ptr (which can be null)
+    std::visit(
+        [&](auto* next_ptr) -> void {
+            if (next_ptr == nullptr) return;  // state was destructed first
+            next_ptr->prev_ptr_ = prev_ptr_;
+        },
+        next_ptr_
+    );
+}
+
+DiffCheckpoint::DiffCheckpoint(CheckpointableState& state, ssize_t drop) :
+    LinkedListCheckpoint(state), updates_(), drop_(drop) {}
+
+DiffCheckpoint::DiffCheckpoint(CheckpointableState& state, std::span<const Update> diff) :
+    LinkedListCheckpoint(state), updates_(), drop_(diff.size()) {
+    if (auto* prev_ptr = static_cast<DiffCheckpoint*>(prev_ptr_)) {
+        prev_ptr->commit_updates(std::vector<Update>(diff.begin(), diff.end()));
+        assert(prev_ptr->drop_ == 0);
+    }
+}
+
+DiffCheckpoint::~DiffCheckpoint() {
+    // If we're not the oldest checkpoint, we need to transfer our information
+    // over so it's not lost
+    if (auto* prev_ptr = static_cast<DiffCheckpoint*>(prev_ptr_)) {
+        assert(prev_ptr->drop_ == 0);
+        for (auto& updates : updates_) prev_ptr->commit_updates(std::move(updates));
+        prev_ptr->drop_ = drop_;
+    }
+}
+
+void DiffCheckpoint::commit_updates(std::vector<Update> updates) {
+    assert(0 <= drop_ and static_cast<size_t>(drop_) <= updates.size());
+
+    if (drop_) {
+        updates.erase(updates.begin(), updates.begin() + drop_);
+        drop_ = 0;
+    }
+
+    updates_.emplace_back(std::move(updates));
+}
+
+void DiffCheckpoint::revert_updates(std::vector<Update> updates) {
+    assert(0 <= drop_ and static_cast<size_t>(drop_) <= updates.size());
+
+    if (not drop_) return;  // nothing to do
+
+    // We want to track the updates that would revert the changes from the
+    // current state.
+    // In C++23 we could use assign_range() which would be nicer
+    auto relevant = std::move(updates) | std::views::take(drop_) | std::views::reverse |
+                    std::views::transform([](const Update& up) { return up.inverse(); });
+    updates_.emplace_back(relevant.begin(), relevant.end());
+
+    drop_ = 0;
+}
+
+}  // namespace dwave::optimization