From 891ab8bac543d82a23cb0fd1f6d48b85f5c3acf1 Mon Sep 17 00:00:00 2001
From: Nikolaj Bjorner <nbjorner@microsoft.com>
Date: Sun, 15 Oct 2023 12:37:14 -0700
Subject: [PATCH] #6523

fixup looping
---
 src/sat/smt/arith_solver.cpp | 53 ++++++++++++++++++++++++++++--------
 src/sat/smt/arith_solver.h   |  4 ++-
 src/smt/theory_lra.cpp       |  1 -
 3 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/src/sat/smt/arith_solver.cpp b/src/sat/smt/arith_solver.cpp
index 17ab03ee6..2921e5018 100644
--- a/src/sat/smt/arith_solver.cpp
+++ b/src/sat/smt/arith_solver.cpp
@@ -83,6 +83,7 @@ namespace arith {
 
         m_new_eq = false;
         flush_bound_axioms();
+        propagate_nla();
 
         unsigned qhead = m_asserted_qhead;
         while (m_asserted_qhead < m_asserted.size() && !s().inconsistent() && m.inc()) {
@@ -301,6 +302,22 @@ namespace arith {
         m_explanation.add_pair(j, v);
     }
 
+    void solver::add_equality(lpvar j, rational const& k, lp::explanation const& exp) {
+        TRACE("arith", tout << "equality " << j << " " << k << "\n");
+        theory_var v;
+        if (k == 1)
+            v = m_one_var;
+        else if (k == 0)
+            v = m_zero_var;
+        else if (!m_value2var.find(k, v))
+            return;
+        theory_var w = lp().local_to_external(j);
+        if (w < 0)
+            return;
+        lpvar i = register_theory_var_in_lar_solver(v);
+        add_eq(i, j, exp, true);
+    }
+
     bool solver::add_eq(lpvar u, lpvar v, lp::explanation const& e, bool is_fixed) {
         if (s().inconsistent())
             return false;
@@ -1414,14 +1431,6 @@ namespace arith {
             core.push_back(~mk_ineq_literal(ineq));
         set_conflict_or_lemma(hint_type::nla_h, core, false);
     }
-
-    void solver::assume_literals() {
-        for (auto const& ineq : m_nla->literals()) {
-            auto lit = mk_ineq_literal(ineq);
-            ctx.mark_relevant(lit);
-            s().set_phase(lit);
-        }
-    }
     
     sat::literal solver::mk_ineq_literal(nla::ineq const& ineq) {
         bool is_lower = true, sign = true, is_eq = false;
@@ -1462,9 +1471,7 @@ namespace arith {
         lbool r = m_nla->check();
         switch (r) {
         case l_false:
-            assume_literals();
-            for (const nla::lemma& l : m_nla->lemmas())
-                false_case_of_check_nla(l);
+            add_lemmas();
             break;
         case l_true:
             if (assume_eqs())
@@ -1476,6 +1483,30 @@ namespace arith {
         return r;
     }
 
+    void solver::add_lemmas() {
+        for (auto const& ineq : m_nla->literals()) {
+            auto lit = mk_ineq_literal(ineq);
+            ctx.mark_relevant(lit);
+            s().set_phase(lit);
+        }
+        for (const nla::lemma& l : m_nla->lemmas())
+            false_case_of_check_nla(l);
+        if (!propagate_eqs()) 
+            return;
+        for (auto const& [v,k,e] : m_nla->fixed_equalities())
+            add_equality(v, k, e);
+        for (auto const& [i,j,e] : m_nla->equalities())
+            add_eq(i,j,e,false);
+    }
+
+    void solver::propagate_nla() {
+        if (m_nla) {
+            m_nla->propagate();
+            add_lemmas();
+            lp().collect_more_rows_for_lp_propagation();
+        }
+    }
+
     void solver::get_antecedents(literal l, sat::ext_justification_idx idx, literal_vector& r, bool probing) {
         auto& jst = euf::th_explain::from_index(idx);
         ctx.get_th_antecedents(l, jst, r, probing);
diff --git a/src/sat/smt/arith_solver.h b/src/sat/smt/arith_solver.h
index 801eb474c..b3d10ab52 100644
--- a/src/sat/smt/arith_solver.h
+++ b/src/sat/smt/arith_solver.h
@@ -407,6 +407,9 @@ namespace arith {
         bool  check_delayed_eqs();
         lbool check_lia();
         lbool check_nla();
+        void add_lemmas();
+        void propagate_nla();
+        void add_equality(lpvar v, rational const& k, lp::explanation const& exp);
         bool is_infeasible() const;
 
         nlsat::anum const& nl_value(theory_var v, scoped_anum& r) const;
@@ -463,7 +466,6 @@ namespace arith {
         void set_evidence(lp::constraint_index idx);
         void assign(literal lit, literal_vector const& core, svector<enode_pair> const& eqs, euf::th_proof_hint const* pma);
 
-        void assume_literals();
         sat::literal mk_ineq_literal(nla::ineq const& ineq);
         void false_case_of_check_nla(const nla::lemma& l);        
         void dbg_finalize_model(model& mdl);
diff --git a/src/smt/theory_lra.cpp b/src/smt/theory_lra.cpp
index abbcea199..9e8f57e2c 100644
--- a/src/smt/theory_lra.cpp
+++ b/src/smt/theory_lra.cpp
@@ -2173,7 +2173,6 @@ public:
     }
 
     void add_equality(lpvar j, rational const& k, lp::explanation const& exp) {
-        //verbose_stream() << "equality " << j << " " << k << "\n";
         TRACE("arith", tout << "equality " << j << " " << k << "\n");
         theory_var v;
         if (k == 1)