perf(stn): Use theard-local storage for bound propagation persistent …

…memory
plaans · Nov 21, 2024 · 5d585a1 · 5d585a1
1 parent fb8457e
commit 5d585a1
Showing 1 changed file with 19 additions and 13 deletions.
diff --git a/solver/src/reasoners/stn/theory/bound_propagation.rs b/solver/src/reasoners/stn/theory/bound_propagation.rs
@@ -1,4 +1,4 @@
-use std::cmp::Reverse;
+use std::{cell::RefCell, cmp::Reverse};
 
 use crate::{
     collections::{heap::IdxHeap, ref_store::RefMap},
@@ -10,6 +10,10 @@ use super::{
     IntCst, PropagatorId, SignedVar, StnTheory, INT_CST_MAX,
 };
 
+thread_local! {
+    static STATE: RefCell<Dij> = RefCell::new(Dij::default())
+}
+
 /// Process all bound changes in `stn.pending_bound_changes` and propagate them by:
 ///  - updating the bounds of all active edges in the network
 ///  - activating/deactivating all edges that are entailed/diabled by the current bounds
@@ -23,20 +27,22 @@ pub fn process_bound_changes(
     doms: &mut Domains,
     cycle_detection: impl Fn(SignedVar) -> bool,
 ) -> Result<(), InvalidUpdate> {
-    let mut dij = Dij::default();
-    dij.clear();
+    // acquire our thread-local working memory to avoid repeated allocations
+    STATE.with_borrow_mut(|dij| {
+        dij.clear();
 
-    for update in &stn.pending_bound_changes {
-        dij.add_modified_bound(
-            update.var,
-            update.previous_ub,
-            update.new_ub,
-            update.is_from_bound_propagation,
-        );
-    }
-    stn.pending_bound_changes.clear();
+        for update in &stn.pending_bound_changes {
+            dij.add_modified_bound(
+                update.var,
+                update.previous_ub,
+                update.new_ub,
+                update.is_from_bound_propagation,
+            );
+        }
+        stn.pending_bound_changes.clear();
 
-    dij.run(stn, doms, cycle_detection)
+        dij.run(stn, doms, cycle_detection)
+    })
 }
 
 #[derive(Default, Clone)]