Add unification based projection function

Heuristic for EMBPR is to unify terms that occur in uninterpreted contents. Walk partitions that are pure f(t) and an impure f(s)
and attempt to unify t with s ensuring that merges from s preserves satisfiability.

src/qe/mbp/mbp_euf.cpp

@@ -7,15 +7,12 @@ Copyright (c) 2025 Microsoft Corporation
 #include "ast/ast_util.h"
 #include "ast/for_each_expr.h"
 #include "qe/mbp/mbp_euf.h"
-#include "qe/mbp/mbp_term_graph.h"
 
 namespace mbp {
-    euf_project_plugin::euf_project_plugin(ast_manager& m): project_plugin(m) {
-        
+    euf_project_plugin::euf_project_plugin(ast_manager& m): project_plugin(m) {        
     }
 
     euf_project_plugin::~euf_project_plugin() {
-
     }
     
     bool euf_project_plugin::project1(model& model, app* var, app_ref_vector& vars, expr_ref_vector& lits) {
@@ -26,9 +23,6 @@ namespace mbp {
         return basic_family_id;        
     }
     
-    bool euf_project_plugin::operator()(model& model, app_ref_vector& vars, expr_ref_vector& lits) {
-        return false;
-    }
 
     void euf_project_plugin::solve_eqs(model& model, app_ref_vector& vars, expr_ref_vector& lits, vector<def>& defs) {
         flatten_and(lits);
@@ -51,10 +45,8 @@ namespace mbp {
             }
         }
     }
-    
-    bool euf_project_plugin::project(model& model, app_ref_vector& vars, expr_ref_vector& lits, vector<def>& defs) {
-        if (vars.empty())
-            return false;
+
+    bool euf_project_plugin::solve_eqs_saturate(model& model, app_ref_vector& vars, expr_ref_vector& lits, vector<def>& defs) {
         unsigned sz = defs.size();
         while (true) {
             unsigned sz1 = defs.size();
@@ -62,14 +54,51 @@ namespace mbp {
             if (sz1 == defs.size())
                 break;
         }
+        return sz < defs.size();
+    }
 
-        if (sz < defs.size())
-            return true;
-        
+
+    bool euf_project_plugin::operator()(model& model, app_ref_vector& vars, expr_ref_vector& lits) {
+        if (vars.empty())
+            return false;
         // check if there is a variable of uninterp sort
         if (all_of(vars, [&](expr* v) { return !m.is_uninterp(v->get_sort()); }))
             return false;
 
+        vector<def> defs;
+        if (solve_eqs_saturate(model, vars, lits, defs))
+            return true;
+        term_graph tg(m);
+        tg.add_lits(lits);
+        for (auto v : vars) 
+            if (m.is_uninterp(v->get_sort()))
+                tg.add_var(v);
+
+        auto result = tg.project(model);
+        lits.reset();
+        lits.append(result);
+
+        unsigned j = 0;
+        for (app* v : vars) 
+            if (!m.is_uninterp(v->get_sort())) 
+                vars[j++] = v;
+        vars.shrink(j);
+
+        return true;
+    }
+
+    bool euf_project_plugin::project(model& model, app_ref_vector& vars, expr_ref_vector& lits, vector<def>& defs) {
+        if (vars.empty())
+            return false;
+
+        // check if there is a variable of uninterp sort
+        if (all_of(vars, [&](expr* v) { return !m.is_uninterp(v->get_sort()); }))
+            return false;
+
+        if (solve_eqs_saturate(model, vars, lits, defs))
+            return true;
+        
+
         term_graph tg(m);
         tg.add_lits(lits);
         for (auto v : vars) 
@@ -78,19 +107,36 @@ namespace mbp {
 
         expr_ref_vector terms(m);
         for (expr* f : subterms::all(lits))
-            if (!m.is_bool(f))
-                terms.push_back(f);
-
+            terms.push_back(f);
 
         // try to project on representatives:
         tg.add_model_based_terms(model, terms);
-        unsigned j = 0;        
-        for (unsigned i = 0; i < vars.size(); ++i) {
-            app* v = vars.get(i);
+
+        // save rep_of. It gets reset by tg.get_partition.
+        struct scoped_reset {
+            euf_project_plugin& p;
+            scoped_reset(euf_project_plugin& p): p(p) {}
+            ~scoped_reset() { p.m_reps.reset(); p.m_parents.reset(); }
+        };
+        scoped_reset _reset(*this);
+        expr_ref_vector pinned(m);
+        for (auto t : terms) {
+            m_reps.insert(t, tg.rep_of(t));
+            pinned.push_back(tg.rep_of(t));
+            if (is_app(t)) {
+                for (expr* arg : *to_app(t))
+                    m_parents.insert_if_not_there(arg, {}).push_back(t);
+            }
+        }
+        for (auto [key, value] : m_reps)
+            verbose_stream() << mk_pp(key, m) << " -> " << mk_pp(value, m) << "\n";
+        unsigned j = 0;
+        bool solved = false;
+        for (app* v : vars) {
             if (m.is_uninterp(v->get_sort())) {
-                expr* r = tg.rep_of(v);
+                expr* r = m_reps[v];
                 if (r) 
-                    defs.push_back({ expr_ref(v, m), expr_ref(r, m) });
+                    defs.push_back({ expr_ref(v, m), expr_ref(r, m) }), solved = true;
                 else
                     vars[j++] = v;
             }
@@ -99,18 +145,131 @@ namespace mbp {
         }
         vars.shrink(j);
 
-        if (sz < defs.size())
+        if (solved)
             return true;
-        
+       
 
         // walk equivalence classes.
         // try to unify with pure classes.
+        // some parent of y contains f(y) and class of f contains f(t) with repr(t) != nullptr
+        // merge y = t, provided parent f(y) is not distinct from f(t)
+        // unification can break distinctness, and should require a full solver call.
+        // a conservative rule is to walk only the equivalence class of f(y) to check if it
+        // contains f(t),
+        // or recursively if there is a term g(f(t)) that is equal to g(f(y))
+        // thus, take equivalence classes that contain a repr and unify
+        //
+        // this can still break satisfiability if f(y) is required distinct from f(t) even
+        // if g(f(t)) = g(f(y))
+        // so as a simplification unification is only allowed if g(f(y)) is the only parent of f(y) 
 
-
-        // try to build a fresh term using available signature
-
+        auto Ps = tg.get_partition(model);
+        for (auto const& p : Ps) {
+            expr* r = m_reps[p[0]];
+            if (!r)
+                continue;
+            for (auto e : p) {
+                if (!is_app(e))
+                    continue;
+                app* a = to_app(e);
+                func_decl* f = a->get_decl();
+                if (!is_uninterp(f))
+                    continue;
+                if (all_of(*a, [&](expr* arg) { return !!m_reps[arg]; }))
+                    continue;
+                if (try_unify(tg, a, p, vars, defs)) 
+                    return true;
+            }
+        }        
         
         return false;
     }
 
+    bool euf_project_plugin::try_unify(term_graph& tg, app* a, expr_ref_vector const& partition, app_ref_vector& vars, vector<def>& defs) {
+
+        auto same_decl = [&](expr* x, expr* y) {
+            if (!is_app(x) || !is_app(y))
+                return false;
+            app* a = to_app(x);
+            app* b = to_app(y);
+            if (a->get_decl() != b->get_decl())
+                return false;
+            return a->get_num_args() == b->get_num_args();
+        };
+
+        auto is_var = [&](expr* x) {
+            return is_uninterp_const(x) && m.is_uninterp(x->get_sort()) && vars.contains(to_app(x));
+        };
+        
+        for (auto e : partition) {
+            if (a == e)
+                continue;
+            verbose_stream() << "Unify " << mk_pp(a, m) << " with " << mk_pp(e, m) << "\n";
+            if (!same_decl(a, e))
+                continue;
+            app* b = to_app(e);
+            if (!all_of(*b, [&](expr* arg) { return !!m_reps[arg]; }))
+                continue;
+            // same function symbol. All kids of b are defined, some kid of a is not defined.
+            svector<std::pair<expr*, expr*>> todo;
+            obj_map<expr, expr*> soln;
+            for (unsigned i = 0; i < b->get_num_args(); ++i) {
+                expr* x = a->get_arg(i);
+                expr* y = b->get_arg(i);
+                todo.push_back({x, y});
+            }            
+            while (!todo.empty()) {
+                auto [x, y] = todo.back();
+                todo.pop_back();
+                auto rx = m_reps[x];
+                auto ry = m_reps[y];
+                SASSERT(ry);
+                if (rx == ry) 
+                    continue;
+                if (rx)
+                    break; // fail
+                if (m_parents[x].size() > 1)
+                    // a crude safey hatch to preserve models
+                    // we could require that every parent of x
+                    // has a pure representative which comes from a correspondnig
+                    // parent of y
+                    break;
+                expr* s = nullptr;
+                if (soln.find(x, s)) {
+                    if (s != ry)
+                        break;
+                    continue;
+                }
+                if (is_var(x)) {
+                    soln.insert(x, ry);
+                    continue;
+                }
+                if (!same_decl(x, y))
+                    break;
+                app* ax = to_app(x);
+                app* ay = to_app(y);
+                for (unsigned i = 0; i < ax->get_num_args(); ++i) 
+                    todo.push_back({ ax->get_arg(i), ay->get_arg(i) });
+            }
+
+            TRACE("qe", tout << "unification attempt\n";
+                  for (auto [a, b] : todo)                 
+                      tout << mk_pp(a, m) << " == " << mk_pp(b, m) << "\n";
+                  for (auto [key, value] : soln)
+                      tout << mk_pp(key, m) << " -> " << mk_pp(value, m) << "\n";
+                  );
+
+
+            if (todo.empty() && !soln.empty()) {
+                for (auto [key, value] : soln) {
+                    vars.erase(to_app(key));
+                    defs.push_back( { expr_ref(key, m), expr_ref(value, m) });
+                }
+                return true;
+            }
+        }
+        return false;
+    }
+
+
 }

src/qe/mbp/mbp_euf.h

@@ -9,11 +9,16 @@ Copyright (c) 2025 Microsoft Corporation
 
 #include "model/model.h"
 #include "qe/mbp/mbp_plugin.h"
+#include "qe/mbp/mbp_term_graph.h"
 
 namespace mbp {
 
     class euf_project_plugin : public project_plugin {
+        obj_map<expr, expr*> m_reps;
+        obj_map<expr, ptr_vector<expr>> m_parents;
         void solve_eqs(model& model, app_ref_vector& vars, expr_ref_vector& lits, vector<def>& defs);
+        bool solve_eqs_saturate(model& model, app_ref_vector& vars, expr_ref_vector& lits, vector<def>& defs);
+        bool try_unify(term_graph& g, app* a, expr_ref_vector const& partitions, app_ref_vector& vars, vector<def>& defs);
     public:
         euf_project_plugin(ast_manager& m);
         ~euf_project_plugin() override;

src/qe/mbp/mbp_term_graph.cpp

@@ -167,7 +167,7 @@ class term {
     };
     class_props m_class_props;
 
-  public:
+public:
     term(expr_ref const &v, u_map<term *> &app2term)
         : m_expr(v), m_root(this), m_repr(nullptr), m_next(this), m_mark(false),
           m_mark2(false), m_interpreted(false),
@@ -190,7 +190,7 @@ class term {
     class parents {
         term const &t;
 
-      public:
+    public:
         parents(term const &_t) : t(_t) {}
         parents(term const *_t) : t(*_t) {}
         ptr_vector<term>::const_iterator begin() const {
@@ -204,7 +204,7 @@ class term {
     class children {
         term const &t;
 
-      public:
+    public:
         children(term const &_t) : t(_t) {}
         children(term const *_t) : t(*_t) {}
         ptr_vector<term>::const_iterator begin() const {
@@ -577,8 +577,9 @@ term *term_graph::internalize_term(expr *t) {
     }
     merge_flush();
     SASSERT(res);
-    if (m.is_not(t) && is_app(to_app(t)->get_arg(0)) && is_partial_eq(to_app(to_app(t)->get_arg(0)))) {
-        term* p = get_term(to_app(t)->get_arg(0));
+    expr* arg = nullptr;
+    if (m.is_not(t, arg) && is_app(arg) && is_partial_eq(to_app(arg))) {
+        term* p = get_term(arg);
         SASSERT(p);
         p->set_npeq_child();
     }
@@ -879,8 +880,7 @@ void term_graph::pick_repr_class(term *t) {
 void term_graph::pick_repr() {
     // invalidates cache
     m_term2app.reset();
-    DEBUG_CODE(for (term *t
-                    : m_terms)
+    DEBUG_CODE(for (term *t : m_terms)
                    SASSERT(t->deg() == 0 || !t->all_children_ground() ||
                            t->is_cgr()););
     for (term *t : m_terms) t->reset_repr();
@@ -1090,25 +1090,27 @@ class term_graph::projector {
     expr_ref_vector m_pinned; // tracks expr in the maps
 
     expr *mk_pure(term const &t) {
-        TRACE("qe", t.display(tout););
+        TRACE("qe", t.display(tout));
         expr *e = nullptr;
-        if (find_term2app(t, e)) return e;
+        if (find_term2app(t, e))
+            return e;
         e = t.get_expr();
-        if (!is_app(e)) return nullptr;
+        if (!is_app(e))
+            return nullptr;
         app *a = ::to_app(e);
         expr_ref_buffer kids(m);
         for (term *ch : term::children(t)) {
             // prefer a node that resembles current child,
             // otherwise, pick a root representative, if present.
-            if (find_term2app(*ch, e)) { kids.push_back(e); }
-            else if (m_root2rep.find(ch->get_root().get_id(), e)) {
+            if (find_term2app(*ch, e))
+                kids.push_back(e); 
+            else if (m_root2rep.find(ch->get_root().get_id(), e)) 
                 kids.push_back(e);
-            }
-            else { return nullptr; }
+            else
+                return nullptr; 
             TRACE("qe_verbose", tout << *ch << " -> " << mk_pp(e, m) << "\n";);
         }
-        expr_ref pure =
-            m_rewriter.mk_app(a->get_decl(), kids.size(), kids.data());
+        expr_ref pure = m_rewriter.mk_app(a->get_decl(), kids.size(), kids.data());
         m_pinned.push_back(pure);
         add_term2app(t, pure);
         return pure;
@@ -1164,7 +1166,7 @@ class term_graph::projector {
         m_tg.reset_marks();
     }
 
-    bool find_app(term &t, expr *&res) {
+    bool find_app(term const &t, expr *&res) {
         return find_term2app(t, res) ||
                m_root2rep.find(t.get_root().get_id(), res);
     }
@@ -1177,9 +1179,9 @@ class term_graph::projector {
 
     void mk_lits(expr_ref_vector &res) {
         expr *e = nullptr;
-        for (auto *lit : m_tg.m_lits) {
-            if (!m.is_eq(lit) && find_app(lit, e)) res.push_back(e);
-        }
+        for (auto *lit : m_tg.m_lits) 
+            if (!m.is_eq(lit) && find_app(lit, e))
+                res.push_back(e);
         TRACE("qe", tout << "literals: " << res << "\n";);
     }
 
@@ -1219,12 +1221,12 @@ class term_graph::projector {
     }
 
     void remove_duplicates(expr_ref_vector &v) {
-        obj_hashtable<expr> seen;
+        expr_fast_mark1 seen;
         unsigned j = 0;
         for (expr *e : v) {
-            if (!seen.contains(e)) {
+            if (!seen.is_marked(e)) {
                 v[j++] = e;
-                seen.insert(e);
+                seen.mark(e);
             }
         }
         v.shrink(j);
@@ -1425,7 +1427,7 @@ class term_graph::projector {
         return out;
     }
 
-  public:
+public:
     projector(term_graph &tg)
         : m_tg(tg), m(m_tg.m), m_rewriter(m), m_pinned(m) {}
 
@@ -1439,9 +1441,9 @@ class term_graph::projector {
         return m_term2app.find(t.get_id(), r);
     }
 
-    expr *find_term2app(term const &t) {
+    expr* rep_of(term const &t) {
         expr *r = nullptr;
-        find_term2app(t, r);
+        find_app(t, r);
         return r;
     }
 
@@ -1709,11 +1711,11 @@ void term_graph::add_model_based_terms(model &mdl,
     m_projector->purify();
 }
 
-expr *term_graph::rep_of(expr *e) {
+expr* term_graph::rep_of(expr *e) {
     SASSERT(m_projector);
     term *t = get_term(e);
     SASSERT(t && "only get representatives");
-    return m_projector->find_term2app(*t);
+    return m_projector->rep_of(*t);    
 }
 
 expr_ref_vector term_graph::dcert(model &mdl, expr_ref_vector const &lits) {
@@ -1834,15 +1836,14 @@ void term_graph::cground_percolate_up(term *t) {
 }
 
 void term_graph::cground_percolate_up(ptr_vector<term> &todo) {
-    term *t;
-
     while (!todo.empty()) {
-        t = todo.back();
+        term* t = todo.back();
         todo.pop_back();
         t->set_cgr(true);
         t->set_class_gr(true);
         for (auto p : term::parents(t->get_root()))
-            if (!p->is_cgr() && p->all_children_ground()) todo.push_back(p);
+            if (!p->is_cgr() && p->all_children_ground())
+                todo.push_back(p);
     }
 }
 
@@ -1852,12 +1853,11 @@ void term_graph::compute_cground() {
         t->set_class_gr(false);
     }
     ptr_vector<term> todo;
-    for (auto t : m_terms) {
-        if (t->is_gr()) { todo.push_back(t); }
-    }
+    for (auto t : m_terms) 
+        if (t->is_gr())
+            todo.push_back(t); 
     cground_percolate_up(todo);
-    DEBUG_CODE(for (auto t
-                    : m_terms) {
+    DEBUG_CODE(for (auto t : m_terms) {
         bool isclsg = true;
         for (auto c : term::children(t)) isclsg &= c->is_class_gr();
         SASSERT(t->deg() == 0 || !isclsg || t->is_cgr());