From 3f3d05856787dd6a6ae030a81800381c310a5969 Mon Sep 17 00:00:00 2001
From: Nikolaj Bjorner <nbjorner@microsoft.com>
Date: Mon, 7 Feb 2022 06:15:49 +0200
Subject: [PATCH] extract also units from search state

---
 src/smt/smt_context.cpp | 16 ++++++++++++++++
 src/smt/smt_context.h   |  2 ++
 src/smt/smt_kernel.cpp  |  6 +++++-
 src/smt/smt_kernel.h    |  6 ++++++
 src/smt/smt_solver.cpp  |  4 ++++
 src/solver/solver.cpp   |  1 +
 src/solver/solver.h     |  2 ++
 7 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/src/smt/smt_context.cpp b/src/smt/smt_context.cpp
index 60eccf962..735da24eb 100644
--- a/src/smt/smt_context.cpp
+++ b/src/smt/smt_context.cpp
@@ -4613,6 +4613,22 @@ namespace smt {
         return result;
     }
 
+    void context::get_units(expr_ref_vector& result) {
+        expr_mark visited;
+        for (expr* fml : result)
+            visited.mark(fml);
+        for (literal lit : m_assigned_literals) {
+            if (get_assign_level(lit) > m_base_lvl)
+                break;
+            expr_ref e(m);
+            literal2expr(lit, e);
+            if (visited.is_marked(e))
+                continue;
+            result.push_back(std::move(e));
+        }
+    }
+
+
     failure context::get_last_search_failure() const {
         return m_last_search_failure;
     }
diff --git a/src/smt/smt_context.h b/src/smt/smt_context.h
index 89b94412d..b8c8a7ee9 100644
--- a/src/smt/smt_context.h
+++ b/src/smt/smt_context.h
@@ -1691,6 +1691,8 @@ namespace smt {
 
         void get_assertions(ptr_vector<expr> & result) { m_asserted_formulas.get_assertions(result); }
 
+        void get_units(expr_ref_vector& result);
+
         /*
          * user-propagator
          */
diff --git a/src/smt/smt_kernel.cpp b/src/smt/smt_kernel.cpp
index 4a2c8bae9..347260c84 100644
--- a/src/smt/smt_kernel.cpp
+++ b/src/smt/smt_kernel.cpp
@@ -177,6 +177,10 @@ namespace smt {
     void kernel::get_assignments(expr_ref_vector & result) {
         m_imp->m_kernel.get_assignments(result);
     }
+
+    void kernel::get_units(expr_ref_vector & result) {
+        m_imp->m_kernel.get_units(result);
+    }    
         
     void kernel::get_relevant_labels(expr * cnstr, buffer<symbol> & result) {
         m_imp->m_kernel.get_relevant_labels(cnstr, result);
@@ -280,4 +284,4 @@ namespace smt {
         m_imp->m_kernel.user_propagate_register_created(r);
     }
 
-};
\ No newline at end of file
+};
diff --git a/src/smt/smt_kernel.h b/src/smt/smt_kernel.h
index 681bdd55b..28680e7a6 100644
--- a/src/smt/smt_kernel.h
+++ b/src/smt/smt_kernel.h
@@ -202,6 +202,12 @@ namespace smt {
            \brief Return the set of formulas assigned by the kernel.
         */
         void get_assignments(expr_ref_vector & result);
+
+
+        /**
+           \brief Return units assigned by the kernel.
+        */
+        void get_units(expr_ref_vector& result);
         
         /**
            \brief Return the set of relevant labels in the last check command.
diff --git a/src/smt/smt_solver.cpp b/src/smt/smt_solver.cpp
index 03f55585a..1cb9b26e1 100644
--- a/src/smt/smt_solver.cpp
+++ b/src/smt/smt_solver.cpp
@@ -318,6 +318,10 @@ namespace {
             return m_context.get_formula(idx);
         }
 
+        void get_units_core(expr_ref_vector& units) override {
+            m_context.get_units(units);
+        }
+
         expr_ref_vector cube(expr_ref_vector& vars, unsigned cutoff) override {
             ast_manager& m = get_manager();
             if (!m_cuber) {
diff --git a/src/solver/solver.cpp b/src/solver/solver.cpp
index 7dcca2ecc..d14648058 100644
--- a/src/solver/solver.cpp
+++ b/src/solver/solver.cpp
@@ -242,6 +242,7 @@ expr_ref_vector solver::get_units() {
     ast_manager& m = get_manager();
     expr_ref_vector fmls(m), result(m), tmp(m);
     get_assertions(fmls);
+    get_units_core(fmls);
     obj_map<expr, bool> units;
     for (expr* f : fmls) {
         if (m.is_not(f, f) && is_literal(m, f)) {
diff --git a/src/solver/solver.h b/src/solver/solver.h
index 090cd9d9b..13692c857 100644
--- a/src/solver/solver.h
+++ b/src/solver/solver.h
@@ -261,6 +261,8 @@ public:
     */
     expr_ref_vector get_units();
 
+    virtual void get_units_core(expr_ref_vector& units) {} 
+
     expr_ref_vector get_non_units();
 
     virtual expr_ref_vector get_trail() = 0; // { return expr_ref_vector(get_manager()); }