From 6787d87623fba607d18ad79184d509c3d52f6885 Mon Sep 17 00:00:00 2001
From: Nikolaj Bjorner <nbjorner@microsoft.com>
Date: Mon, 6 Jan 2025 13:16:07 -0800
Subject: [PATCH] hoist update stack creation

Signed-off-by: Nikolaj Bjorner <nbjorner@microsoft.com>
---
 src/ast/sls/sls_bv_lookahead.cpp | 68 +++++++++++++++++++-------------
 src/ast/sls/sls_bv_lookahead.h   |  3 ++
 2 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/src/ast/sls/sls_bv_lookahead.cpp b/src/ast/sls/sls_bv_lookahead.cpp
index 9c7522c53..8c12c8506 100644
--- a/src/ast/sls/sls_bv_lookahead.cpp
+++ b/src/ast/sls/sls_bv_lookahead.cpp
@@ -490,18 +490,10 @@ namespace sls {
 
         // TRACE("bv_verbose", tout << "lookahead update " << mk_bounded_pp(t, m) << " := " << new_value << "\n";);
 
-        insert_update_stack(t);
-
-        unsigned max_depth = get_depth(t);
-        for (unsigned depth = max_depth; depth <= max_depth; ++depth) {
+        for (unsigned depth = m_min_depth; depth <= m_max_depth; ++depth) {
             for (unsigned i = 0; i < m_update_stack[depth].size(); ++i) {
                 auto a = m_update_stack[depth][i];
                 TRACE("bv_verbose", tout << "update " << mk_bounded_pp(a, m) << " depth: " << depth  << "\n";);
-                for (auto p : ctx.parents(a)) {
-                    insert_update_stack(p);
-                    max_depth = std::max(max_depth, get_depth(p));
-                }
-               
                 if (t != a) {
                     if (bv.is_bv(a))
                         m_ev.eval(a);
@@ -512,9 +504,7 @@ namespace sls {
                     score += get_weight(a) * (new_score(a) - old_score(a));
                 }
             }
-            m_update_stack[depth].reset();
         }
-        m_in_update_stack.reset();
         restore_lookahead();
         m_ev.restore_bool_values(restore_point);
 
@@ -523,6 +513,26 @@ namespace sls {
         return score;
     }
 
+    void bv_lookahead::populate_update_stack(expr* t) {
+        insert_update_stack(t);
+        m_min_depth = m_max_depth = get_depth(t);
+        for (unsigned depth = m_max_depth; depth <= m_max_depth; ++depth) {
+            for (unsigned i = 0; i < m_update_stack[depth].size(); ++i) {
+                auto a = m_update_stack[depth][i];
+                for (auto p : ctx.parents(a)) {
+                    insert_update_stack(p);
+                    m_max_depth = std::max(m_max_depth, get_depth(p));
+                }
+            }
+        }
+    }
+
+    void bv_lookahead::clear_update_stack() {
+        for (unsigned i = m_min_depth; i <= m_max_depth; ++i)
+            m_update_stack[i].reset();
+        m_in_update_stack.reset();
+    }
+
     void bv_lookahead::try_set(expr* u, bvect const& new_value) {
         if (!wval(u).can_set(new_value))
             return;
@@ -551,7 +561,9 @@ namespace sls {
 
     void bv_lookahead::add_updates(expr* u) {
         if (m.is_bool(u)) {
+            populate_update_stack(u);
             try_flip(u);
+            clear_update_stack();
             return;
         }
         SASSERT(bv.is_bv(u));
@@ -567,6 +579,7 @@ namespace sls {
         v.bits().copy_to(v.nw, m_v_saved);
         m_v_saved.copy_to(v.nw, m_v_updated);
 
+        populate_update_stack(u);
         // flip a single bit
         for (unsigned i = 0; i < v.bw && i < 32 ; ++i) {
             m_v_updated.set(i, !m_v_updated.get(i));
@@ -579,25 +592,26 @@ namespace sls {
             try_set(u, m_v_updated);
             m_v_updated.set(j, !m_v_updated.get(j));
         }
-        if (v.bw <= 1)
-            return;
+        if (v.bw > 1) {
 
-        // increment
-        m_v_saved.copy_to(v.nw, m_v_updated);
-        v.add1(m_v_updated);
-        try_set(u, m_v_updated);
+            // increment
+            m_v_saved.copy_to(v.nw, m_v_updated);
+            v.add1(m_v_updated);
+            try_set(u, m_v_updated);
 
-        // decrement
-        m_v_saved.copy_to(v.nw, m_v_updated);
-        v.sub1(m_v_updated);
-        try_set(u, m_v_updated);
+            // decrement
+            m_v_saved.copy_to(v.nw, m_v_updated);
+            v.sub1(m_v_updated);
+            try_set(u, m_v_updated);
 
-        // invert
-        m_v_saved.copy_to(v.nw, m_v_updated);
-        for (unsigned i = 0; i < v.nw; ++i)
-            m_v_updated[i] = ~m_v_updated[i];
-        v.clear_overflow_bits(m_v_updated);
-        try_set(u, m_v_updated);
+            // invert
+            m_v_saved.copy_to(v.nw, m_v_updated);
+            for (unsigned i = 0; i < v.nw; ++i)
+                m_v_updated[i] = ~m_v_updated[i];
+            v.clear_overflow_bits(m_v_updated);
+            try_set(u, m_v_updated);
+        }
+        clear_update_stack();
     }
 
     /**
diff --git a/src/ast/sls/sls_bv_lookahead.h b/src/ast/sls/sls_bv_lookahead.h
index 0849a1769..e707da867 100644
--- a/src/ast/sls/sls_bv_lookahead.h
+++ b/src/ast/sls/sls_bv_lookahead.h
@@ -90,6 +90,9 @@ namespace sls {
 
         bv_valuation& wval(expr* e) const;
 
+        unsigned m_max_depth = 0, m_min_depth = 0;
+        void populate_update_stack(expr* e);
+        void clear_update_stack();
         void insert_update_stack(expr* e);
         void insert_update(expr* e);        
         void restore_lookahead();