From 6e72a97727073f7b1f37c37aa2b7b45a5797c944 Mon Sep 17 00:00:00 2001
From: Jakob Rath <git@jakobrath.eu>
Date: Mon, 21 Nov 2022 17:25:15 +0100
Subject: [PATCH] Refactor assignment and search state

---
 src/math/polysat/CMakeLists.txt           |   1 +
 src/math/polysat/assignment.cpp           | 107 ++++++++++++++++++++
 src/math/polysat/assignment.h             | 114 ++++++++++++++++++++++
 src/math/polysat/boolean.h                |   3 -
 src/math/polysat/constraint.cpp           |   4 +
 src/math/polysat/constraint.h             |  12 +--
 src/math/polysat/op_constraint.cpp        |   8 +-
 src/math/polysat/op_constraint.h          |   3 +-
 src/math/polysat/search_state.cpp         |  33 ++-----
 src/math/polysat/search_state.h           |  96 +++---------------
 src/math/polysat/simplify_clause.cpp      |   6 +-
 src/math/polysat/smul_fl_constraint.h     |   3 +-
 src/math/polysat/solver.cpp               |  69 +++++++------
 src/math/polysat/solver.h                 |  15 +--
 src/math/polysat/superposition.cpp        |   6 +-
 src/math/polysat/types.h                  |   3 +
 src/math/polysat/ule_constraint.cpp       |  24 ++---
 src/math/polysat/ule_constraint.h         |   3 +-
 src/math/polysat/umul_ovfl_constraint.cpp |   4 +-
 src/math/polysat/umul_ovfl_constraint.h   |   3 +-
 src/math/polysat/variable_elimination.cpp |  18 ++--
 src/math/polysat/variable_elimination.h   |   2 +-
 22 files changed, 326 insertions(+), 211 deletions(-)
 create mode 100644 src/math/polysat/assignment.cpp
 create mode 100644 src/math/polysat/assignment.h

diff --git a/src/math/polysat/CMakeLists.txt b/src/math/polysat/CMakeLists.txt
index b7fc66903..5d4d3f3f5 100644
--- a/src/math/polysat/CMakeLists.txt
+++ b/src/math/polysat/CMakeLists.txt
@@ -1,5 +1,6 @@
 z3_add_component(polysat
   SOURCES
+    assignment.cpp
     boolean.cpp
     clause.cpp
     clause_builder.cpp
diff --git a/src/math/polysat/assignment.cpp b/src/math/polysat/assignment.cpp
new file mode 100644
index 000000000..70bbb55ed
--- /dev/null
+++ b/src/math/polysat/assignment.cpp
@@ -0,0 +1,107 @@
+/*++
+Copyright (c) 2021 Microsoft Corporation
+
+Module Name:
+
+    polysat substitution and assignment
+
+Author:
+
+    Nikolaj Bjorner (nbjorner) 2021-03-19
+    Jakob Rath 2021-04-06
+
+--*/
+
+#include "math/polysat/assignment.h"
+#include "math/polysat/solver.h"
+
+namespace polysat {
+
+    substitution::substitution(pdd p)
+        : m_subst(std::move(p)) { }
+
+    substitution::substitution(dd::pdd_manager& m)
+        : m_subst(m.one()) { }
+
+    substitution substitution::add(pvar var, rational const& value) const {
+        return {m_subst.subst_add(var, value)};
+    }
+
+    pdd substitution::apply_to(pdd const& p) const {
+        return p.subst_val(m_subst);
+    }
+
+    bool substitution::value(pvar var, rational& out_value) const {
+        return m_subst.subst_get(var, out_value);
+    }
+
+    assignment::assignment(solver& s)
+        : m_solver(&s) { }
+
+    assignment assignment::clone() const {
+        assignment a(s());
+        a.m_pairs = m_pairs;
+        a.m_subst.reserve(m_subst.size());
+        for (unsigned i = m_subst.size(); i-- > 0; )
+            if (m_subst[i])
+                a.m_subst.set(i, alloc(substitution, *m_subst[i]));
+        a.m_subst_trail = m_subst_trail;
+        return a;
+    }
+
+    substitution& assignment::subst(unsigned sz) {
+        return const_cast<substitution&>(std::as_const(*this).subst(sz));
+    }
+
+    substitution const& assignment::subst(unsigned sz) const {
+        m_subst.reserve(sz + 1);
+        if (!m_subst[sz])
+            m_subst.set(sz, alloc(substitution, s().sz2pdd(sz)));
+        return *m_subst[sz];
+    }
+
+    void assignment::push(pvar var, rational const& value) {
+        SASSERT(all_of(m_pairs, [var](assignment_item_t const& item) { return item.first != var; }));
+        m_pairs.push_back({var, value});
+        unsigned const sz = s().size(var);
+        substitution& sub = subst(sz);
+        m_subst_trail.push_back(sub);
+        sub = sub.add(var, value);
+        SASSERT_EQ(sub, *m_subst[sz]);
+    }
+
+    void assignment::pop() {
+        substitution& sub = m_subst_trail.back();
+        unsigned sz = sub.bit_width();
+        SASSERT_EQ(sz, s().size(m_pairs.back().first));
+        *m_subst[sz] = sub;
+        m_subst_trail.pop_back();
+        m_pairs.pop_back();
+    }
+
+    pdd assignment::apply_to(pdd const& p) const {
+        unsigned const sz = p.power_of_2();
+        return subst(sz).apply_to(p);
+    }
+
+    std::ostream& substitution::display(std::ostream& out) const {
+        char const* delim = "";
+        pdd p = m_subst;
+        while (!p.is_val()) {
+            SASSERT(p.lo().is_val());
+            out << delim << "v" << p.var() << " := " << p.lo();
+            delim = " ";
+            p = p.hi();
+        }
+        return out;
+    }
+
+    std::ostream& assignment::display(std::ostream& out) const {
+        char const* delim = "";
+        for (auto const& [var, value] : m_pairs) {
+            out << delim << assignment_pp(s(), var, value);
+            delim = " ";
+        }
+        return out;
+    }
+}
diff --git a/src/math/polysat/assignment.h b/src/math/polysat/assignment.h
new file mode 100644
index 000000000..db2e52b68
--- /dev/null
+++ b/src/math/polysat/assignment.h
@@ -0,0 +1,114 @@
+/*++
+Copyright (c) 2021 Microsoft Corporation
+
+Module Name:
+
+    polysat substitution and assignment
+
+Author:
+
+    Nikolaj Bjorner (nbjorner) 2021-03-19
+    Jakob Rath 2021-04-06
+
+--*/
+#pragma once
+#include "math/polysat/types.h"
+
+namespace polysat {
+
+    using assignment_item_t = std::pair<pvar, rational>;
+
+    class substitution_iterator {
+        pdd m_current;
+        substitution_iterator(pdd current) : m_current(std::move(current)) {}
+        friend class substitution;
+
+    public:
+        using value_type = assignment_item_t;
+        using difference_type = std::ptrdiff_t;
+        using pointer = value_type const*;
+        using reference = value_type const&;
+        using iterator_category = std::input_iterator_tag;
+
+        substitution_iterator& operator++() {
+            SASSERT(!m_current.is_val());
+            m_current = m_current.hi();
+            return *this;
+        }
+
+        value_type operator*() const {
+            SASSERT(!m_current.is_val());
+            return { m_current.var(), m_current.lo().val() };
+        }
+
+        bool operator==(substitution_iterator const& other) const { return m_current == other.m_current; }
+        bool operator!=(substitution_iterator const& other) const { return !operator==(other); }
+    };
+
+    /** Substitution for a single bit width. */
+    class substitution {
+        pdd m_subst;
+
+        substitution(pdd p);
+
+    public:
+        substitution(dd::pdd_manager& m);
+        substitution add(pvar var, rational const& value) const;
+        pdd apply_to(pdd const& p) const;
+
+        bool value(pvar var, rational& out_value) const;
+
+        bool empty() const { return m_subst.is_one(); }
+
+        pdd const& to_pdd() const { return m_subst; }
+        unsigned bit_width() const { return to_pdd().power_of_2(); }
+
+        bool operator==(substitution const& other) const { return &m_subst.manager() == &other.m_subst.manager() && m_subst == other.m_subst; }
+        bool operator!=(substitution const& other) const { return !operator==(other); }
+
+        std::ostream& display(std::ostream& out) const;
+
+        using const_iterator = substitution_iterator;
+        const_iterator begin() const { return {m_subst}; }
+        const_iterator end() const { return {m_subst.manager().one()}; }
+    };
+
+    /** Full variable assignment, may include variables of varying bit widths. */
+    class assignment {
+        solver*                                 m_solver;
+        vector<assignment_item_t>               m_pairs;  // TODO: do we still need this?
+        mutable scoped_ptr_vector<substitution> m_subst;
+        vector<substitution>                    m_subst_trail;
+
+        substitution& subst(unsigned sz);
+        solver& s() const { return *m_solver; }
+    public:
+        assignment(solver& s);
+        // prevent implicit copy, use clone() if you do need a copy
+        assignment(assignment const&) = delete;
+        assignment& operator=(assignment const&) = delete;
+        assignment(assignment&&) = default;
+        assignment& operator=(assignment&&) = default;
+        assignment clone() const;
+
+        void push(pvar var, rational const& value);
+        void pop();
+
+        pdd apply_to(pdd const& p) const;
+
+        bool empty() const { return pairs().empty(); }
+        substitution const& subst(unsigned sz) const;
+        vector<assignment_item_t> const& pairs() const { return m_pairs; }
+        using const_iterator = decltype(m_pairs)::const_iterator;
+        const_iterator begin() const { return pairs().begin(); }
+        const_iterator end() const { return pairs().end(); }
+
+        std::ostream& display(std::ostream& out) const;
+    };
+
+    using assignment_t = assignment;
+
+    inline std::ostream& operator<<(std::ostream& out, substitution const& sub) { return sub.display(out); }
+
+    inline std::ostream& operator<<(std::ostream& out, assignment const& a) { return a.display(out); }
+}
diff --git a/src/math/polysat/boolean.h b/src/math/polysat/boolean.h
index b908d9773..948a2d278 100644
--- a/src/math/polysat/boolean.h
+++ b/src/math/polysat/boolean.h
@@ -12,12 +12,9 @@ Author:
 --*/
 #pragma once
 #include "math/polysat/types.h"
-#include "util/sat_literal.h"
 
 namespace polysat {
 
-    class clause;
-
     class bool_var_manager {
 
         enum class kind_t {
diff --git a/src/math/polysat/constraint.cpp b/src/math/polysat/constraint.cpp
index 9cc8fe34e..89e2dd10b 100644
--- a/src/math/polysat/constraint.cpp
+++ b/src/math/polysat/constraint.cpp
@@ -24,6 +24,10 @@ Author:
 
 namespace polysat {
 
+    bool constraint::is_currently_false(solver& s, bool is_positive) const {
+        return is_currently_false(s.assignment(), is_positive);
+    }
+
     bool signed_constraint::is_eq() const {
         return is_positive() && m_constraint->is_eq();
     }
diff --git a/src/math/polysat/constraint.h b/src/math/polysat/constraint.h
index 03330c619..9b2b47b2a 100644
--- a/src/math/polysat/constraint.h
+++ b/src/math/polysat/constraint.h
@@ -15,7 +15,7 @@ Author:
 #include "math/polysat/boolean.h"
 #include "math/polysat/types.h"
 #include "math/polysat/interval.h"
-#include "math/polysat/search_state.h"
+#include "math/polysat/assignment.h"
 #include "math/polysat/univariate/univariate_solver.h"
 #include <iostream>
 
@@ -85,11 +85,11 @@ namespace polysat {
         virtual std::ostream& display(std::ostream& out) const = 0;
 
         virtual bool is_always_false(bool is_positive) const = 0;
-        virtual bool is_currently_false(solver& s, bool is_positive) const = 0;
-        virtual bool is_currently_false(solver& s, assignment_t const& sub, bool is_positive) const = 0;
+        bool is_currently_false(solver& s, bool is_positive) const;
+        virtual bool is_currently_false(assignment const& a, bool is_positive) const = 0;
         bool is_always_true(bool is_positive) const { return is_always_false(!is_positive); }
         bool is_currently_true(solver& s, bool is_positive) const { return is_currently_false(s, !is_positive); }
-        bool is_currently_true(solver& s, assignment_t const& sub, bool is_positive) const { return is_currently_false(s, sub, !is_positive); }
+        bool is_currently_true(assignment const& a, bool is_positive) const { return is_currently_false(a, !is_positive); }
 
         virtual void narrow(solver& s, bool is_positive, bool first) = 0;
         virtual inequality as_inequality(bool is_positive) const = 0;
@@ -155,8 +155,8 @@ namespace polysat {
         bool is_always_true() const { return get()->is_always_false(is_negative()); }
         bool is_currently_false(solver& s) const { return get()->is_currently_false(s, is_positive()); }
         bool is_currently_true(solver& s) const { return get()->is_currently_false(s, is_negative()); }
-        bool is_currently_false(solver& s, assignment_t const& sub) const { return get()->is_currently_false(s, sub, is_positive()); }
-        bool is_currently_true(solver& s, assignment_t const& sub) const { return get()->is_currently_false(s, sub, is_negative()); }
+        bool is_currently_false(assignment const& a) const { return get()->is_currently_false(a, is_positive()); }
+        bool is_currently_true(assignment const& a) const { return get()->is_currently_false(a, is_negative()); }
         lbool bvalue(solver& s) const;
         void narrow(solver& s, bool first) { get()->narrow(s, is_positive(), first); }
         inequality as_inequality() const { return get()->as_inequality(is_positive()); }
diff --git a/src/math/polysat/op_constraint.cpp b/src/math/polysat/op_constraint.cpp
index 2638ef1f5..61bf41d1f 100644
--- a/src/math/polysat/op_constraint.cpp
+++ b/src/math/polysat/op_constraint.cpp
@@ -80,12 +80,8 @@ namespace polysat {
         return is_always_false(is_positive, p(), q(), r());
     }
 
-    bool op_constraint::is_currently_false(solver& s, bool is_positive) const {
-        return is_always_false(is_positive, s.subst(p()), s.subst(q()), s.subst(r()));
-    }
-
-    bool op_constraint::is_currently_false(solver& s, assignment_t const& sub, bool is_positive) const {
-        return is_always_false(is_positive, s.subst(sub, p()), s.subst(sub, q()), s.subst(sub, r()));
+    bool op_constraint::is_currently_false(assignment const& a, bool is_positive) const {
+        return is_always_false(is_positive, a.apply_to(p()), a.apply_to(q()), a.apply_to(r()));
     }
 
     std::ostream& op_constraint::display(std::ostream& out, lbool status) const {
diff --git a/src/math/polysat/op_constraint.h b/src/math/polysat/op_constraint.h
index 0f84d72a7..777db0f35 100644
--- a/src/math/polysat/op_constraint.h
+++ b/src/math/polysat/op_constraint.h
@@ -60,8 +60,7 @@ namespace polysat {
         std::ostream& display(std::ostream& out, lbool status) const override;
         std::ostream& display(std::ostream& out) const override;
         bool is_always_false(bool is_positive) const override;
-        bool is_currently_false(solver& s, bool is_positive) const override;
-        bool is_currently_false(solver& s, assignment_t const& sub, bool is_positive) const override;
+        bool is_currently_false(assignment const& a, bool is_positive) const override;
         void narrow(solver& s, bool is_positive, bool first) override;
         inequality as_inequality(bool is_positive) const override { throw default_exception("is not an inequality"); }
         unsigned hash() const override;
diff --git a/src/math/polysat/search_state.cpp b/src/math/polysat/search_state.cpp
index d7966a026..6eb8e9131 100644
--- a/src/math/polysat/search_state.cpp
+++ b/src/math/polysat/search_state.cpp
@@ -8,7 +8,7 @@ Module Name:
 Author:
 
     Nikolaj Bjorner (nbjorner) 2021-03-19
-    Jakob Rath 2021-04-6
+    Jakob Rath 2021-04-06
 
 --*/
 
@@ -89,7 +89,7 @@ namespace polysat {
     std::ostream& search_state::display(std::ostream& out) const {
         for (auto const& item : m_items)
             display(item, out) << " ";
-        return out;        
+        return out;
     }
 
     bool search_state::value(pvar v, rational& val) const {
@@ -103,28 +103,7 @@ namespace polysat {
 
     void search_state::push_assignment(pvar p, rational const& r) {
         m_items.push_back(search_item::assignment(p));
-        m_assignment.push_back({p, r});
-        unsigned sz = s.size(p);
-        pdd& s = assignment(sz);
-        m_subst_trail.push_back(s);
-        s = s.subst_add(p, r);
-        SASSERT(s == *m_subst[sz]);
-    }
-
-    pdd& search_state::assignment(unsigned sz) const {
-        m_subst.reserve(sz + 1);
-        if (!m_subst[sz])
-            m_subst.set(sz, alloc(pdd, s.sz2pdd(sz).one()));
-        return *m_subst[sz];
-    }
-
-    void search_state::pop_assignment() {
-        m_assignment.pop_back();
-        pdd& s = m_subst_trail.back();
-        auto& m = s.manager();
-        unsigned sz = m.power_of_2();
-        *m_subst[sz] = s;
-        m_subst_trail.pop_back();
+        m_assignment.push(p, r);
     }
 
     void search_state::push_boolean(sat::literal lit) {
@@ -133,8 +112,10 @@ namespace polysat {
 
     void search_state::pop() {
         auto const& item = m_items.back();
-        if (item.is_assignment() && !m_assignment.empty() && item.var() == m_assignment.back().first) 
-            pop_assignment();
+        if (item.is_assignment()) {
+            SASSERT_EQ(item.var(), m_assignment.pairs().back().first);
+            m_assignment.pop();
+        }
         m_items.pop_back();
     }
 
diff --git a/src/math/polysat/search_state.h b/src/math/polysat/search_state.h
index 0958e7964..544f304f3 100644
--- a/src/math/polysat/search_state.h
+++ b/src/math/polysat/search_state.h
@@ -8,20 +8,15 @@ Module Name:
 Author:
 
     Nikolaj Bjorner (nbjorner) 2021-03-19
-    Jakob Rath 2021-04-6
+    Jakob Rath 2021-04-06
 
 --*/
 #pragma once
-#include "math/polysat/boolean.h"
 #include "math/polysat/types.h"
+#include "math/polysat/assignment.h"
 
 namespace polysat {
 
-    typedef std::pair<pvar, rational> assignment_item_t;
-    typedef vector<assignment_item_t> assignment_t;
-
-    class solver;
-
     enum class search_item_k
     {
         assignment,
@@ -46,7 +41,7 @@ namespace polysat {
         bool is_resolved() const { return m_resolved; }
         search_item_k kind() const { return m_kind; }
         pvar var() const { SASSERT(is_assignment()); return m_var; }
-        sat::literal lit() const { SASSERT(is_boolean()); return m_lit; }        
+        sat::literal lit() const { SASSERT(is_boolean()); return m_lit; }
         void set_resolved() { m_resolved = true; }
     };
 
@@ -54,27 +49,28 @@ namespace polysat {
         solver&             s;
 
         vector<search_item> m_items;
-        assignment_t        m_assignment;  // First-order part of the search state
-        mutable scoped_ptr_vector<pdd>         m_subst;
-        vector<pdd>         m_subst_trail;
+        assignment          m_assignment;
 
         bool value(pvar v, rational& r) const;
 
     public:
-        search_state(solver& s): s(s) {}
+        search_state(solver& s): s(s), m_assignment(s) {}
         unsigned size() const { return m_items.size(); }
         search_item const& back() const { return m_items.back(); }
         search_item const& operator[](unsigned i) const { return m_items[i]; }
 
         assignment_t const& assignment() const { return m_assignment; }
-        pdd& assignment(unsigned sz) const;        
+        substitution const& subst(unsigned sz) const { return m_assignment.subst(sz); }
+
+        // TODO: implement the following method if we actually need the assignments without resolved items already during conflict resolution
+        //       (no separate trail needed, just a second m_subst and an index into the trail, I think)
+        //       (update on set_resolved? might be one iteration too early, looking at the old solver::resolve_conflict loop)
+        substitution const& unresolved_assignment(unsigned sz) const;
 
         void push_assignment(pvar p, rational const& r);
         void push_boolean(sat::literal lit);
         void pop();
 
-        void pop_assignment();
-
         void set_resolved(unsigned i) { m_items[i].set_resolved(); }
 
         using const_iterator = decltype(m_items)::const_iterator;
@@ -106,74 +102,4 @@ namespace polysat {
 
     inline std::ostream& operator<<(std::ostream& out, search_item_pp const& p) { return p.verbose ? p.s.display_verbose(p.i, out) : p.s.display(p.i, out); }
 
-    // Go backwards over the search_state.
-    // If new entries are added during processing an item, they will be queued for processing next after the current item.
-    class search_iterator {
-
-        search_state*   m_search;
-
-        unsigned current;
-        unsigned first;  // highest index + 1
-
-        struct idx_range {
-            unsigned current;
-            unsigned first;  // highest index + 1
-        };
-        vector<idx_range>     m_index_stack;
-
-        void init() {
-            first = m_search->size();
-            current = first;  // we start one before the beginning
-        }
-
-        void try_push_block() {
-            if (first != m_search->size()) {
-                m_index_stack.push_back({current, first});
-                init();
-            }
-        }
-
-        void pop_block() {
-            current = m_index_stack.back().current;
-            // We don't restore 'first', otherwise 'next()' will immediately push a new block again.
-            // Instead, the current block is merged with the popped one.
-            m_index_stack.pop_back();
-        }
-
-        unsigned last() {
-            return m_index_stack.empty() ? 0 : m_index_stack.back().first;
-        }
-
-    public:
-        search_iterator(search_state& search):
-            m_search(&search) {
-            init();
-        }
-
-        void set_resolved() {
-            m_search->set_resolved(current);
-        }
-
-        search_item const& operator*() {
-            return (*m_search)[current];
-        }
-
-        bool next() {
-#if 1  // If you want to resolve over constraints that have been added during conflict resolution, enable this.
-            try_push_block();
-#endif
-            if (current > last()) {
-                --current;
-                return true;
-            }
-            else {
-                SASSERT(current == last());
-                if (m_index_stack.empty())
-                    return false;
-                pop_block();
-                return next();
-            }
-        }
-    };
-
 }
diff --git a/src/math/polysat/simplify_clause.cpp b/src/math/polysat/simplify_clause.cpp
index 86ed73ac5..b7feac925 100644
--- a/src/math/polysat/simplify_clause.cpp
+++ b/src/math/polysat/simplify_clause.cpp
@@ -133,9 +133,9 @@ namespace polysat {
                 cl[j++] = cl[i];
             else {
                 DEBUG_CODE({
-                    auto a = s.assignment();
-                    a.push_back({v, k});
-                    SASSERT(s.lit2cnstr(lit).is_currently_false(s, a));
+                    auto a = s.assignment().clone();
+                    a.push(v, k);
+                    SASSERT(s.lit2cnstr(lit).is_currently_false(a));
                 });
             }
         }
diff --git a/src/math/polysat/smul_fl_constraint.h b/src/math/polysat/smul_fl_constraint.h
index dffd0be6b..f7388a2b6 100644
--- a/src/math/polysat/smul_fl_constraint.h
+++ b/src/math/polysat/smul_fl_constraint.h
@@ -36,8 +36,7 @@ namespace polysat {
         std::ostream& display(std::ostream& out) const override;
         bool is_always_false(bool is_positive) const override { return false; }
         void narrow(solver& s, bool is_positive, bool first) override;
-        bool is_currently_false(solver & s, bool is_positive) const override { return false; }
-        bool is_currently_false(solver& s, assignment_t const& sub, bool is_positive) const override { return false; }
+        bool is_currently_false(assignment const& a, bool is_positive) const override { return false; }
 
         inequality as_inequality(bool is_positive) const override { throw default_exception("is not an inequality"); }
         unsigned hash() const override;
diff --git a/src/math/polysat/solver.cpp b/src/math/polysat/solver.cpp
index 615c9ec39..9b088b0c7 100644
--- a/src/math/polysat/solver.cpp
+++ b/src/math/polysat/solver.cpp
@@ -97,7 +97,7 @@ namespace polysat {
         return *m_pdd[sz];
     }
 
-    dd::pdd_manager& solver::var2pdd(pvar v) {
+    dd::pdd_manager& solver::var2pdd(pvar v) const {
         return sz2pdd(size(v));
     }
 
@@ -693,15 +693,14 @@ namespace polysat {
 
         SASSERT(is_conflict());
 
-        search_iterator search_it(m_search);
-        while (search_it.next()) {
-            auto& item = *search_it;
-            search_it.set_resolved();
+        for (unsigned i = m_search.size(); i-- > 0; ) {
+            auto& item = m_search[i];
+            m_search.set_resolved(i);
             if (item.is_assignment()) {
                 // Resolve over variable assignment
                 pvar v = item.var();
                 if (!m_conflict.is_relevant_pvar(v)) {
-                    m_search.pop_assignment();
+                    // m_search.pop_assignment();
                     continue;
                 }
                 LOG_H2("Working on " << search_item_pp(m_search, item));
@@ -712,7 +711,7 @@ namespace polysat {
                     revert_decision(v);
                     return;
                 }
-                m_search.pop_assignment();
+                // m_search.pop_assignment();
             }
             else {
                 // Resolve over boolean literal
@@ -834,25 +833,17 @@ namespace polysat {
 
     void solver::backjump_and_learn(unsigned jump_level, clause& lemma) {
 #ifndef NDEBUG
-        assignment_t old_assignment;
-        // We can't use solver::assignment() here because we already used search_sate::pop_assignment().
-        // TODO: fix search_state design; it should show a consistent state.
-        search_iterator search_it(m_search);
-        while (search_it.next()) {
-            auto& item = *search_it;
-            if (item.is_assignment()) {
-                pvar v = item.var();
-                old_assignment.push_back({v, get_value(v)});
-            }
-        }
+        polysat::assignment old_assignment = assignment().clone();
         sat::literal_vector lemma_invariant_todo;
         SASSERT(lemma_invariant_part1(lemma, old_assignment, lemma_invariant_todo));
         // SASSERT(lemma_invariant(lemma, old_assignment));
 #endif
         clause_ref_vector side_lemmas = m_conflict.take_side_lemmas();
         sat::literal_vector narrow_queue = m_conflict.take_narrow_queue();
+
         m_conflict.reset();
         backjump(jump_level);
+
         for (sat::literal lit : narrow_queue) {
             switch (m_bvars.value(lit)) {
             case l_true:
@@ -899,10 +890,12 @@ namespace polysat {
         }
     }
 
-    bool solver::lemma_invariant_part1(clause const& lemma, assignment_t const& assignment, sat::literal_vector& out_todo) {
+    bool solver::lemma_invariant_part1(clause const& lemma, polysat::assignment const& a, sat::literal_vector& out_todo) {
         SASSERT(out_todo.empty());
         LOG("Lemma: " << lemma);
-        // LOG("assignment: " << assignment);
+        LOG("assignment: " << a);
+        // TODO: fix
+#if 0
         for (sat::literal lit : lemma) {
             auto const c = lit2cnstr(lit);
             bool const currently_false = c.is_currently_false(*this, assignment);
@@ -912,25 +905,42 @@ namespace polysat {
             else
                 SASSERT(m_bvars.value(lit) == l_false || currently_false);
         }
+#endif
         return true;
     }
 
     bool solver::lemma_invariant_part2(sat::literal_vector const& todo) {
+        LOG("todo: " << todo);
         // Check that undef literals are now propagated by the side lemmas.
+        //
+        // Unfortunately, this isn't always possible.
+        // Consider if the first side lemma contains a constraint that comes from a boolean decision:
+        //
+        //      76: v10 + v7 + -1*v0 + -1 == 0      [ l_true decide@5 pwatched active ]
+        //
+        // When we now backtrack behind the decision level of the literal, then we cannot propagate the side lemma,
+        // and some literals of the main lemma may still be undef at this point.
+        //
+        // So it seems that using constraints from a non-asserting lemma makes the new lemma also non-asserting (if it isn't already).
+#if 1
         for (sat::literal lit : todo)
             SASSERT(m_bvars.value(lit) == l_false);
+#endif
         return true;
     }
 
-    bool solver::lemma_invariant(clause const& lemma, assignment_t const& old_assignment) {
+    bool solver::lemma_invariant(clause const& lemma, polysat::assignment const& old_assignment) {
         LOG("Lemma: " << lemma);
         // LOG("old_assignment: " << old_assignment);
+        // TODO: fix
+#if 0
         for (sat::literal lit : lemma) {
             auto const c = lit2cnstr(lit);
             bool const currently_false = c.is_currently_false(*this, old_assignment);
             LOG("  " << lit_pp(*this, lit) << "    currently_false? " << currently_false);
             SASSERT(m_bvars.value(lit) == l_false || currently_false);
         }
+#endif
         return true;
     }
 
@@ -1167,9 +1177,7 @@ namespace polysat {
     }
 
     std::ostream& assignments_pp::display(std::ostream& out) const {
-        for (auto const& [var, val] : s.assignment())
-            out << assignment_pp(s, var, val) << " ";
-        return out;
+        return out << s.assignment();
     }
 
     std::ostream& assignment_pp::display(std::ostream& out) const {
@@ -1279,18 +1287,7 @@ namespace polysat {
     }
 
     pdd solver::subst(pdd const& p) const {
-        unsigned sz = p.manager().power_of_2();
-        pdd const& s = m_search.assignment(sz);
-        return p.subst_val(s);
-    }
-
-    pdd solver::subst(assignment_t const& sub, pdd const& p) const {
-        unsigned sz = p.manager().power_of_2();
-        pdd s = p.manager().mk_val(1);
-        for (auto const& [var, val] : sub)
-            if (size(var) == sz)
-                s = p.manager().subst_add(s, var, val);
-        return p.subst_val(s);
+        return assignment().apply_to(p);
     }
 
     /** Check that boolean assignment and constraint evaluation are consistent */
diff --git a/src/math/polysat/solver.h b/src/math/polysat/solver.h
index d95139e70..e1ed2a54c 100644
--- a/src/math/polysat/solver.h
+++ b/src/math/polysat/solver.h
@@ -31,6 +31,7 @@ Author:
 #include "math/polysat/justification.h"
 #include "math/polysat/linear_solver.h"
 #include "math/polysat/search_state.h"
+#include "math/polysat/assignment.h"
 #include "math/polysat/trail.h"
 #include "math/polysat/viable.h"
 #include "math/polysat/log.h"
@@ -58,6 +59,7 @@ namespace polysat {
             stats() { reset(); }
         };
 
+        friend class assignment;
         friend class constraint;
         friend class ule_constraint;
         friend class umul_ovfl_constraint;
@@ -125,8 +127,6 @@ namespace polysat {
         unsigned_vector          m_size;     // store size of variables (bit width)
 
         search_state             m_search;
-        assignment_t const& assignment() const { return m_search.assignment(); }
-        pdd subst(assignment_t const& sub, pdd const& p) const;
 
         unsigned                 m_qhead = 0; // next item to propagate (index into m_search)
         unsigned                 m_level = 0;
@@ -150,7 +150,6 @@ namespace polysat {
             m_qhead_trail.pop_back();
         }
 
-
         unsigned size(pvar v) const { return m_size[v]; }
 
         /**
@@ -161,7 +160,9 @@ namespace polysat {
         void del_var();
 
         dd::pdd_manager& sz2pdd(unsigned sz) const;
-        dd::pdd_manager& var2pdd(pvar v);
+        dd::pdd_manager& var2pdd(pvar v) const;
+
+        assignment const& assignment() const { return m_search.assignment(); }
 
         void push_level();
         void pop_levels(unsigned num_levels);
@@ -235,8 +236,8 @@ namespace polysat {
 
         bool invariant();
         static bool invariant(signed_constraints const& cs);
-        bool lemma_invariant(clause const& lemma, assignment_t const& assignment);
-        bool lemma_invariant_part1(clause const& lemma, assignment_t const& assignment, sat::literal_vector& out_todo);
+        bool lemma_invariant(clause const& lemma, polysat::assignment const& a);
+        bool lemma_invariant_part1(clause const& lemma, polysat::assignment const& a, sat::literal_vector& out_todo);
         bool lemma_invariant_part2(sat::literal_vector const& todo);
         bool wlist_invariant();
         bool assignment_invariant();
@@ -429,7 +430,7 @@ namespace polysat {
 
     };  // class solver
 
-    class assignments_pp {
+    class assignments_pp {  // TODO: can probably remove this now.
         solver const& s;
     public:
         assignments_pp(solver const& s): s(s) {}
diff --git a/src/math/polysat/superposition.cpp b/src/math/polysat/superposition.cpp
index d05718429..562e98536 100644
--- a/src/math/polysat/superposition.cpp
+++ b/src/math/polysat/superposition.cpp
@@ -15,8 +15,6 @@ Author:
 #include "math/polysat/log.h"
 #include "math/polysat/solver.h"
 
-// TODO: rename file
-
 namespace polysat {
 
     struct inference_sup : public inference {
@@ -37,8 +35,8 @@ namespace polysat {
         SASSERT(c1.is_currently_true(s));
         SASSERT(c2.is_currently_false(s));
         LOG_H3("Resolving upon v" << v);
-        LOG("c1: " << c1);
-        LOG("c2: " << c2);
+        LOG("c1: " << lit_pp(s, c1));
+        LOG("c2: " << lit_pp(s, c2));
         pdd a = c1.eq();
         pdd b = c2.eq();
         unsigned degree_a = a.degree();
diff --git a/src/math/polysat/types.h b/src/math/polysat/types.h
index e3c08434d..61dac4ea9 100644
--- a/src/math/polysat/types.h
+++ b/src/math/polysat/types.h
@@ -18,6 +18,7 @@ Author:
 #include "util/scoped_ptr_vector.h"
 #include "util/var_queue.h"
 #include "util/ref_vector.h"
+#include "util/sat_literal.h"
 #include "math/dd/dd_pdd.h"
 #include "math/dd/dd_bdd.h"
 #include "math/dd/dd_fdd.h"
@@ -25,6 +26,8 @@ Author:
 namespace polysat {
 
     class solver;
+    class clause;
+
     typedef dd::pdd pdd;
     typedef dd::bdd bdd;
     typedef dd::bddv bddv;
diff --git a/src/math/polysat/ule_constraint.cpp b/src/math/polysat/ule_constraint.cpp
index f2cde7c90..b5041f240 100644
--- a/src/math/polysat/ule_constraint.cpp
+++ b/src/math/polysat/ule_constraint.cpp
@@ -17,7 +17,7 @@ Notes:
 
    Rewrite rules to simplify expressions.
    In the following let k, k1, k2 be values.
-  
+
    - k1 <= k2       ==>   0 <= 0 if k1 <= k2
    - k1 <= k2       ==>   1 <= 0 if k1 >  k2
    - 0 <= p         ==>   0 <= 0
@@ -39,7 +39,7 @@ Notes:
 
 TODO: clause simplifications:
 
-   - p + k <= p  ==> p + k <= k & p != 0   for k != 0  
+   - p + k <= p  ==> p + k <= k & p != 0   for k != 0
    - p*q = 0     ==> p = 0 or q = 0        applies to any factoring
    - 2*p <= 2*q  ==> (p >= 2^n-1 & q < 2^n-1) or (p >= 2^n-1 = q >= 2^n-1 & p <= q)
                  ==> (p >= 2^n-1 => q < 2^n-1 or p <= q) &
@@ -50,7 +50,7 @@ TODO: clause simplifications:
 
 Note:
      case p <= p + k is already covered because we test (lhs - rhs).is_val()
-      
+
      It can be seen as an instance of lemma 5.2 of Supratik and John.
 
 --*/
@@ -168,7 +168,7 @@ namespace polysat {
     void ule_constraint::narrow(solver& s, bool is_positive, bool first) {
         auto p = s.subst(lhs());
         auto q = s.subst(rhs());
-        
+
         signed_constraint sc(this, is_positive);
 
         LOG_H3("Narrowing " << sc);
@@ -215,29 +215,23 @@ namespace polysat {
         return is_always_false(is_positive, lhs(), rhs());
     }
 
-    bool ule_constraint::is_currently_false(solver& s, bool is_positive) const {
-        auto p = s.subst(lhs());
-        auto q = s.subst(rhs());
-        return is_always_false(is_positive, p, q);
-    }
-
-    bool ule_constraint::is_currently_false(solver& s, assignment_t const& sub, bool is_positive) const { 
-        auto p = s.subst(sub, lhs());
-        auto q = s.subst(sub, rhs());
+    bool ule_constraint::is_currently_false(assignment const& a, bool is_positive) const {
+        auto p = a.apply_to(lhs());
+        auto q = a.apply_to(rhs());
         return is_always_false(is_positive, p, q);
     }
 
     inequality ule_constraint::as_inequality(bool is_positive) const {
         if (is_positive)
             return inequality(lhs(), rhs(), false, this);
-        else 
+        else
             return inequality(rhs(), lhs(), true, this);
     }
 
     unsigned ule_constraint::hash() const {
     	return mk_mix(lhs().hash(), rhs().hash(), kind());
     }
-    
+
     bool ule_constraint::operator==(constraint const& other) const {
         return other.is_ule() && lhs() == other.to_ule().lhs() && rhs() == other.to_ule().rhs();
     }
diff --git a/src/math/polysat/ule_constraint.h b/src/math/polysat/ule_constraint.h
index ddbb3cc96..ee09d1265 100644
--- a/src/math/polysat/ule_constraint.h
+++ b/src/math/polysat/ule_constraint.h
@@ -37,8 +37,7 @@ namespace polysat {
         std::ostream& display(std::ostream& out, lbool status) const override;
         std::ostream& display(std::ostream& out) const override;
         bool is_always_false(bool is_positive) const override;
-        bool is_currently_false(solver& s, bool is_positive) const override;
-        bool is_currently_false(solver& s, assignment_t const& sub, bool is_positive) const override;
+        bool is_currently_false(assignment const& a, bool is_positive) const override;
         void narrow(solver& s, bool is_positive, bool first) override;
         inequality as_inequality(bool is_positive) const override;
         unsigned hash() const override;
diff --git a/src/math/polysat/umul_ovfl_constraint.cpp b/src/math/polysat/umul_ovfl_constraint.cpp
index 2f11d23e6..ca802ca9e 100644
--- a/src/math/polysat/umul_ovfl_constraint.cpp
+++ b/src/math/polysat/umul_ovfl_constraint.cpp
@@ -77,8 +77,8 @@ namespace polysat {
         return is_always_false(is_positive, m_p, m_q);
     }
 
-    bool umul_ovfl_constraint::is_currently_false(solver& s, bool is_positive) const {
-        return is_always_false(is_positive, s.subst(p()), s.subst(q()));
+    bool umul_ovfl_constraint::is_currently_false(assignment const& a, bool is_positive) const {
+        return is_always_false(is_positive, a.apply_to(p()), a.apply_to(q()));
     }
 
     void umul_ovfl_constraint::narrow(solver& s, bool is_positive, bool first) {    
diff --git a/src/math/polysat/umul_ovfl_constraint.h b/src/math/polysat/umul_ovfl_constraint.h
index 53771ab4e..e4f52acb7 100644
--- a/src/math/polysat/umul_ovfl_constraint.h
+++ b/src/math/polysat/umul_ovfl_constraint.h
@@ -39,8 +39,7 @@ namespace polysat {
         std::ostream& display(std::ostream& out) const override;
         bool is_always_false(bool is_positive) const override;
         void narrow(solver& s, bool is_positive, bool first) override;
-        bool is_currently_false(solver & s, bool is_positive) const override;
-        bool is_currently_false(solver& s, assignment_t const& sub, bool is_positive) const override { return false; }
+        bool is_currently_false(assignment const& a, bool is_positive) const override;
 
         inequality as_inequality(bool is_positive) const override { throw default_exception("is not an inequality"); }
         unsigned hash() const override;
diff --git a/src/math/polysat/variable_elimination.cpp b/src/math/polysat/variable_elimination.cpp
index 168490618..b7461c8d5 100644
--- a/src/math/polysat/variable_elimination.cpp
+++ b/src/math/polysat/variable_elimination.cpp
@@ -60,14 +60,14 @@ namespace polysat {
         LOG("lc: " << lc);
         LOG("rest: " << rest);
 
-        assignment_t a;
-        pdd const lcs = eval(lc, core, a);
+        substitution sub(m);
+        pdd const lcs = eval(lc, core, sub);
         LOG("lcs: " << lcs);
         pdd lci = m.zero();
         if (!inv(lcs, lci))
             return;
 
-        pdd const rs = s.subst(a, rest);
+        pdd const rs = sub.apply_to(rest);
         pdd const vs = -rs * lci;  // this is the polynomial that computes v
         LOG("vs: " << vs);
         SASSERT(!vs.free_vars().contains(v));
@@ -102,7 +102,7 @@ namespace polysat {
                 continue;
 
             clause_builder cb(s);
-            for (auto [w, wv] : a)
+            for (auto [w, wv] : sub)
                 cb.insert(~s.eq(s.var(w), wv));
             cb.insert(~c);
             cb.insert(~c_target);
@@ -112,19 +112,19 @@ namespace polysat {
     }
 
     // Evaluate p under assignments in the core.
-    pdd free_variable_elimination::eval(pdd const& p, conflict& core, assignment_t& out_assignment) {
+    pdd free_variable_elimination::eval(pdd const& p, conflict& core, substitution& out_sub) {
         // TODO: this should probably be a helper method on conflict.
         // TODO: recognize constraints of the form "v1 == 27" to be used in the assignment?
         //       (but maybe useful evaluations are always part of core.vars() anyway?)
 
-        assignment_t& a = out_assignment;
-        SASSERT(a.empty());
+        substitution& sub = out_sub;
+        SASSERT(sub.empty());
 
         for (auto v : p.free_vars())
             if (core.contains_pvar(v))
-                a.push_back({v, s.get_value(v)});
+                sub.add(v, s.get_value(v));
 
-        pdd q = s.subst(a, p);
+        pdd q = sub.apply_to(p);
 
         // TODO: like in the old conflict::minimize_vars, we can now try to remove unnecessary variables from a.
 
diff --git a/src/math/polysat/variable_elimination.h b/src/math/polysat/variable_elimination.h
index 3e9a99c30..e1efc828c 100644
--- a/src/math/polysat/variable_elimination.h
+++ b/src/math/polysat/variable_elimination.h
@@ -24,7 +24,7 @@ namespace polysat {
         solver& s;
         void find_lemma(pvar v, conflict& core);
         void find_lemma(pvar v, signed_constraint c, conflict& core);
-        pdd eval(pdd const& p, conflict& core, assignment_t& out_assignment);
+        pdd eval(pdd const& p, conflict& core, substitution& out_sub);
         bool inv(pdd const& p, pdd& out_p_inv);
     public:
         free_variable_elimination(solver& s): s(s) {}