From 9f0b3cdc2520bd089fe1ef4ef81a093960a20273 Mon Sep 17 00:00:00 2001
From: Nikolaj Bjorner <nbjorner@microsoft.com>
Date: Tue, 7 Nov 2023 10:54:15 +0100
Subject: [PATCH] Add utility to expand sub-terms

---
 src/math/lp/lar_solver.cpp | 26 ++++++++++++++++++++++----
 src/math/lp/lar_solver.h   |  7 ++++---
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/src/math/lp/lar_solver.cpp b/src/math/lp/lar_solver.cpp
index 5795e166503..dcdc1224978 100644
--- a/src/math/lp/lar_solver.cpp
+++ b/src/math/lp/lar_solver.cpp
@@ -160,9 +160,8 @@ namespace lp {
         if (tv::is_term(j))
             return j;
         unsigned ext_var_or_term = m_var_register.local_to_external(j);
-        if (tv::is_term(ext_var_or_term)) {
+        if (tv::is_term(ext_var_or_term)) 
             j = ext_var_or_term;
-        }
         return j;
     }
 
@@ -2450,8 +2449,7 @@ namespace lp {
     // a_j.first gives the normalised coefficient,
     // a_j.second givis the column
     bool lar_solver::fetch_normalized_term_column(const lar_term& c, std::pair<mpq, lpvar>& a_j) const {
-        TRACE("lar_solver_terms", tout << "looking for term ";
-        print_term_as_indices(c, tout) << "\n";);
+        TRACE("lar_solver_terms", print_term_as_indices(c, tout << "looking for term ") << "\n";);
         lp_assert(c.is_normalized());
         auto it = m_normalized_terms_to_columns.find(c);
         if (it != m_normalized_terms_to_columns.end()) {
@@ -2463,6 +2461,26 @@ namespace lp {
         return false;
     }
 
+    lar_term lar_solver::unfold_nested_subterms(lar_term const& term) {
+        lar_term result;
+        vector<std::pair<lpvar, mpq>> todo;
+        for (auto const & [j,c] : term.coeffs()) 
+            todo.push_back({j, c});
+        while (!todo.empty()) {
+            auto [j, c] = todo.back();
+            todo.pop_back();
+            auto tv = column2tv(j);
+            if (tv.is_term()) {
+                for (auto const& [j, c2] : get_term(tv).coeffs())
+                    todo.push_back({j, c*c2});                
+            }
+            else
+                result.add_monomial(c, j);
+        }
+        return result;
+    }
+
+
     std::pair<constraint_index, constraint_index> lar_solver::add_equality(lpvar j, lpvar k) {
         vector<std::pair<mpq, var_index>> coeffs;
         if (tv::is_term(j))
diff --git a/src/math/lp/lar_solver.h b/src/math/lp/lar_solver.h
index 2bb0397153e..d01a426805a 100644
--- a/src/math/lp/lar_solver.h
+++ b/src/math/lp/lar_solver.h
@@ -205,8 +205,7 @@ class lar_solver : public column_namer {
 
     static void clean_popped_elements_for_heap(unsigned n, lpvar_heap& set);
     static void clean_popped_elements(unsigned n, indexed_uint_set& set);
-    bool maximize_term_on_tableau(const lar_term& term,
-                                  impq& term_max);
+    bool maximize_term_on_tableau(const lar_term& term, impq& term_max);
     bool costs_are_zeros_for_r_solver() const;
     bool reduced_costs_are_zeroes_for_r_solver() const;
     void set_costs_to_zero(const lar_term& term);
@@ -273,7 +272,7 @@ class lar_solver : public column_namer {
     mutable std::unordered_set<mpq> m_set_of_different_singles;
     mutable mpq m_delta;
 
- public:
+public:
     u_dependency* find_improved_bound(lpvar j, bool is_lower, mpq& bound);
 
     std::ostream& print_explanation(
@@ -430,6 +429,8 @@ class lar_solver : public column_namer {
         return get_term(column2tv(to_column_index(j)));
     }
 
+    lar_term unfold_nested_subterms(lar_term const& term);
+
     inline unsigned row_count() const { return A_r().row_count(); }
     bool var_is_registered(var_index vj) const;
     void clear_inf_heap() {