From 1a742ff784f88802891b21d8319e55ac94b73ff4 Mon Sep 17 00:00:00 2001
From: Nikolaj Bjorner <nbjorner@microsoft.com>
Date: Tue, 2 Jan 2024 09:45:13 -0800
Subject: [PATCH] bugfixes

---
 src/sat/smt/polysat/monomials.cpp   | 24 ++++++++++++++----------
 src/sat/smt/polysat/monomials.h     |  2 ++
 src/sat/smt/polysat_internalize.cpp | 13 ++++++++-----
 src/sat/smt/polysat_solver.cpp      |  2 +-
 src/sat/smt/polysat_solver.h        |  2 +-
 5 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/src/sat/smt/polysat/monomials.cpp b/src/sat/smt/polysat/monomials.cpp
index fba1cc0ff..566a837f1 100644
--- a/src/sat/smt/polysat/monomials.cpp
+++ b/src/sat/smt/polysat/monomials.cpp
@@ -337,17 +337,21 @@ namespace polysat {
         return true;
     }
 
+    std::ostream& monomials::monomial::display(std::ostream& out) const {
+        out << var << " := ";
+        char const* sep = "";
+        for (auto p : args) 
+            if (p.is_var())
+                out << sep << p, sep = " * ";
+            else 
+                out << sep << "(" << p << ")", sep = " * ";
+        out << "\n";
+        return out;
+    }
+
     std::ostream& monomials::display(std::ostream& out) const {
-        for (auto const& mon : m_monomials) {
-            out << mon.var << " := ";
-            char const* sep = "";
-            for (auto p : mon.args) 
-                if (p.is_var())
-                    out << sep << p, sep = " * ";
-                else 
-                    out << sep << "(" << p << ")", sep = " * ";
-            out << "\n";
-        }
+        for (auto const& mon : m_monomials) 
+            mon.display(out);
         return out;
     }
 }
diff --git a/src/sat/smt/polysat/monomials.h b/src/sat/smt/polysat/monomials.h
index e45e0a10f..25bf4de96 100644
--- a/src/sat/smt/polysat/monomials.h
+++ b/src/sat/smt/polysat/monomials.h
@@ -36,6 +36,7 @@ namespace polysat {
             rational   val;
             unsigned size() const { return args.size(); }
             unsigned num_bits() const { return args[0].manager().power_of_2(); }
+            std::ostream& display(std::ostream& out) const;
         };
         vector<monomial>   m_monomials;
         pdd_vector         m_tmp;
@@ -74,6 +75,7 @@ namespace polysat {
                
     };
 
+
     inline std::ostream& operator<<(std::ostream& out, monomials const& m) {
         return m.display(out);
     }
diff --git a/src/sat/smt/polysat_internalize.cpp b/src/sat/smt/polysat_internalize.cpp
index b9de1762e..55c78ba6b 100644
--- a/src/sat/smt/polysat_internalize.cpp
+++ b/src/sat/smt/polysat_internalize.cpp
@@ -793,12 +793,13 @@ namespace polysat {
 
     struct solver::undo_add_eq : public trail {
         solver& s;
+        unsigned sz;
         unsigned index;
     public:
-        undo_add_eq(solver& s, unsigned i) : s(s), index(i) {}
+        undo_add_eq(solver& s, unsigned sz, unsigned i) : s(s), sz(sz), index(i) {}
         void undo() override {
             SASSERT(index == s.m_eqs.back().index());
-            s.m_eq2constraint.remove(index);
+            s.m_eq2constraint[sz].remove(index);
             s.m_eqs.pop_back();
         }
     };
@@ -808,7 +809,9 @@ namespace polysat {
         constraint_id idx;
         std::pair<constraint_id, bool> elem;
         bool is_new = true;
-        if (m_eq2constraint.find(r.index(), elem)) {
+        unsigned sz = r.manager().power_of_2();
+        m_eq2constraint.reserve(sz + 1);
+        if (m_eq2constraint[sz].find(r.index(), elem)) {
             if (elem.second == sign)
                 return elem.first;
             is_new = false;
@@ -816,9 +819,9 @@ namespace polysat {
         auto sc = m_core.eq(p, q);       
         idx = m_core.register_constraint(sc, d);
         if (is_new) {
-            m_eq2constraint.insert(r.index(), { idx, sign });
+            m_eq2constraint[sz].insert(r.index(), { idx, sign });
             m_eqs.push_back(r);
-            ctx.push(undo_add_eq(*this, r.index()));
+            ctx.push(undo_add_eq(*this, sz, r.index()));
         }
         return idx;
     }
diff --git a/src/sat/smt/polysat_solver.cpp b/src/sat/smt/polysat_solver.cpp
index b2744da90..e99733aba 100644
--- a/src/sat/smt/polysat_solver.cpp
+++ b/src/sat/smt/polysat_solver.cpp
@@ -217,9 +217,9 @@ namespace polysat {
         m_var_eqs_head++;
         pdd p = var2pdd(v1);
         pdd q = var2pdd(v2);
-        TRACE("bv", tout << ctx.bpp(n) << " == " << ctx.bpp(var2enode(v2)) << "\n");
         auto d = dependency(v1, v2);
         constraint_id id = eq_constraint(p, q, false, d);
+        TRACE("bv", tout << ctx.bpp(n) << " == " << ctx.bpp(var2enode(v2)) << " " << d << "\n");
         m_core.assign_eh(id, false); 
 
     }
diff --git a/src/sat/smt/polysat_solver.h b/src/sat/smt/polysat_solver.h
index a81571733..e49658361 100644
--- a/src/sat/smt/polysat_solver.h
+++ b/src/sat/smt/polysat_solver.h
@@ -199,7 +199,7 @@ namespace polysat {
         void quot_rem(expr* quot, expr* rem, expr* x, expr* y);
 
         vector<pdd> m_eqs;
-        u_map<std::pair<constraint_id, bool>> m_eq2constraint;
+        vector<u_map<std::pair<constraint_id, bool>>> m_eq2constraint;
         struct undo_add_eq;
         constraint_id eq_constraint(pdd p, pdd q, bool sign, dependency d);