Implement state checkpointing for CollectionNode

.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

@@ -151,6 +151,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.

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;
 
+    /// Set the current state to match the one at the time the given checkpoint was created.
+    void assign_from_checkpoint(State& state, checkpoint_type& checkpoint) const;
+    void assign_from_checkpoint(State& state, checkpoint_type&& checkpoint) const;
+
     const double* buff(const State& state) const override;
 
+    /// Get a checkpoint, an IOU that can be used to return the node to its current state.
+    checkpoint_type checkpoint(State& state) const;
+
     void commit(State&) const override;
 
     std::span<const Update> diff(const State& state) const override;

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,8 @@ class NumberNode : public ArrayOutputMixin<ArrayNode>, public DecisionNode {
 
     // NumberNode methods *****************************************************
 
+    std::unique_ptr<NodeStateCheckpoint> checkpoint(State& state) const; 
+
     // 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 +291,10 @@ class IntegerNode : public NumberNode {
 
     // IntegerNode methods ****************************************************
 
+    /// Set the current state to match the one at the time the given checkpoint was created.
+    void assign_from_checkpoint(State& state, checkpoint_type& checkpoint) const;
+    void assign_from_checkpoint(State& state, checkpoint_type&& checkpoint) const;
+
     // 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/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

@@ -195,6 +195,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,93 @@
+// 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 the oldest checkpoint, just let whatever information we're
+    // holding get destructed with us
+    if (prev_ptr_ == nullptr) return;
+
+    // otherwise we need to transfer our info over
+    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

dwave/optimization/src/nodes/_checkpoints.hpp

@@ -0,0 +1,98 @@
+// 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.
+
+#pragma once
+
+#include <span>
+#include <variant>
+#include <vector>
+
+#include "dwave-optimization/array.hpp"
+#include "dwave-optimization/state.hpp"
+
+namespace dwave::optimization {
+
+class LinkedListCheckpoint;
+
+class CheckpointableState {
+ public:
+    CheckpointableState() = default;
+
+    CheckpointableState(const CheckpointableState&) {}  // the checkpoint pointer is not copied
+    CheckpointableState(CheckpointableState&&) = default;
+
+    CheckpointableState& operator=(const CheckpointableState&) = delete;
+    CheckpointableState& operator=(CheckpointableState&&) = default;
+
+    ~CheckpointableState();
+
+ protected:
+    template <std::derived_from<LinkedListCheckpoint> T>
+    T* checkpoint_ptr() {
+        return static_cast<T*>(prev_ptr_);
+    }
+
+ private:  // todo: private?
+    friend LinkedListCheckpoint;
+
+    // The name is a bit confusing, but by making it match LinkedListCheckpoint::prev_ptr_
+    // it makes the implementations of the various visit methods clearer.
+    LinkedListCheckpoint* prev_ptr_ = nullptr;  // Will be nullptr if there are no checkpoints
+};
+
+class LinkedListCheckpoint : public NodeStateCheckpoint {
+ public:
+    LinkedListCheckpoint() = delete;
+    // We're not moveable or copyable because NodeStateCheckpoint is not.
+
+    LinkedListCheckpoint(CheckpointableState& state);
+
+    ~LinkedListCheckpoint() override;
+
+ protected:  // todo: private?
+    friend CheckpointableState;
+
+    LinkedListCheckpoint* prev_ptr_;
+
+    // Is usually not nullptr unless the state has been destructed
+    std::variant<LinkedListCheckpoint*, CheckpointableState*> next_ptr_;
+};
+
+class DiffCheckpoint : public LinkedListCheckpoint {
+ public:
+    DiffCheckpoint(CheckpointableState& state, std::span<const Update> diff);
+
+    ~DiffCheckpoint() override;
+
+    void commit_updates(std::vector<Update> updates);
+
+    auto detach_updates() {
+        auto updates = std::move(updates_) | std::views::join;
+        assert(updates_.empty());
+        return updates;
+    }
+
+    ssize_t& drop() { return drop_; }
+
+    void revert_updates(std::vector<Update> updates);
+
+ protected:
+    DiffCheckpoint(CheckpointableState& state, ssize_t drop);
+
+ private:
+    std::vector<std::vector<Update>> updates_;
+    ssize_t drop_;
+};
+
+}  // namespace dwave::optimization

dwave/optimization/src/nodes/_state.hpp

@@ -91,6 +91,22 @@ class ArrayStateData {
         assert(size_ >= 0 && static_cast<std::size_t>(size_) == buffer.size());
     }
 
+    // Commit the changes and clear the diff by returning the diff buffer.
+    std::vector<Update> commit_and_detach() {
+        std::vector<Update> tmp;
+        std::swap(updates, tmp);
+        // AlexC: we could now do updates.reserve(tmp.size()) under the assumption
+        // that future update buffers will be a similar size. On the other hand,
+        // not doing this provides another meaningful difference to ::commit().
+        // For now, I think it make sense to not but performance testing needed.
+
+        previous_size_ = buffer.size();
+        assert(size_ >= 0 && static_cast<std::size_t>(size_) == buffer.size());
+
+        assert(updates.empty());
+        return tmp;
+    }
+
     std::span<const Update> diff() const noexcept { return updates; }
 
     // Append a new value to the buffer, tracking the addition in the diff
@@ -149,6 +165,29 @@ class ArrayStateData {
         size_ = buffer.size();
     }
 
+    // Commit the changes and clear the diff by returning the diff buffer.
+    std::vector<Update> revert_and_detach() {
+        assert(previous_size_ >= 0);
+        buffer.resize(previous_size_);
+        const ssize_t size = buffer.size();
+        for (const auto& [index, old, _] : updates | std::views::reverse) {
+            assert(index >= 0);
+            if (index >= size) continue;
+            buffer[index] = old;
+        }
+        size_ = buffer.size();
+
+        std::vector<Update> tmp;
+        std::swap(updates, tmp);
+        // AlexC: we could now do updates.reserve(tmp.size()) under the assumption
+        // that future update buffers will be a similar size. On the other hand,
+        // not doing this provides another meaningful difference to ::commit().
+        // For now, I think it make sense to not but performance testing needed.
+
+        assert(updates.empty());
+        return tmp;
+    }
+
     // Set the value at index, tracking the change in the diff.
     // If allow_emplace is true, do an emplace_back iff the index is equal to the current size.
     bool set(ssize_t i, double value, bool allow_emplace = false) {