Fix NLA optimization regression and relax restore_x

- Relax restore_x() to handle backup/current size mismatches: when
  backup is shorter (new columns added), call
  move_non_basic_columns_to_bounds() to find a feasible solution.
- Fix 100x performance regression in nonlinear optimization: save LP
  optimum before check_nla and return it as bound regardless of NLA
  result, so opt_solver::check_bound() can validate via full re-solve
  with accumulated NLA lemmas.
- Refactor theory_lra::maximize() into three helpers: max_with_lp(),
  max_with_nl(), and max_result().
- Add mk_gt(theory_var, impq const&) overload for building blockers
  from saved LP optimum values.
- Add BNH multi-objective optimization test (7/7 sat in <1s vs 1/7
  in 30s before fix).
- Add restore_x test for backup size mismatch handling.

Fixes #8890

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

src/math/lp/lar_core_solver.h

@@ -81,10 +81,15 @@ public:
     void backup_x() { m_backup_x = m_r_x; }
 
     void restore_x() {
-        SASSERT(m_backup_x.size() == m_r_A.column_count());
-        m_r_x = m_backup_x;
+        unsigned n = m_r_A.column_count();
+        unsigned backup_sz = m_backup_x.size();
+        unsigned copy_sz = std::min(backup_sz, n);
+        for (unsigned j = 0; j < copy_sz; j++)
+            m_r_x[j] = m_backup_x[j];
     }
 
+    unsigned backup_x_size() const { return m_backup_x.size(); }
+
     vector<impq> const& r_x() const { return m_r_x; }
     impq& r_x(unsigned j) { return m_r_x[j]; }
     impq const& r_x(unsigned j) const { return m_r_x[j]; }

src/math/lp/lar_solver.cpp

@@ -467,6 +467,8 @@ namespace lp {
         return ret;
     }
 
+
+
     lp_status lar_solver::solve() {
         if (m_imp->m_status == lp_status::INFEASIBLE || m_imp->m_status == lp_status::CANCELLED)
             return m_imp->m_status;
@@ -2303,16 +2305,6 @@ namespace lp {
         return m_imp->m_constraints.add_term_constraint(j, m_imp->m_columns[j].term(), kind, rs);
     }
 
-    struct lar_solver::scoped_backup {
-        lar_solver& m_s;
-        scoped_backup(lar_solver& s) : m_s(s) {
-            m_s.get_core_solver().backup_x();
-        }
-        ~scoped_backup() {
-            m_s.get_core_solver().restore_x();
-        }
-    };
-   
     void lar_solver::update_column_type_and_bound_with_ub(unsigned j, lp::lconstraint_kind kind, const mpq& right_side, u_dependency* dep) {
         SASSERT(column_has_upper_bound(j));
         if (column_has_lower_bound(j)) {

src/math/lp/lar_solver.h

@@ -72,7 +72,6 @@ class lar_solver : public column_namer {
 
     void clear_columns_with_changed_bounds();
 
-    struct scoped_backup;
  public:
     const indexed_uint_set& columns_with_changed_bounds() const;
     void insert_to_columns_with_changed_bounds(unsigned j);
@@ -437,7 +436,25 @@ public:
     statistics& stats();
 
     void backup_x() { get_core_solver().backup_x(); }
-    void restore_x() { get_core_solver().restore_x(); }
+    void restore_x() {
+        auto& cs = get_core_solver();
+        unsigned backup_sz = cs.backup_x_size();
+        unsigned current_sz = cs.m_n();
+        CTRACE(lar_solver_restore, backup_sz != current_sz,
+               tout << "restore_x: backup_sz=" << backup_sz
+               << " current_sz=" << current_sz << "\n";);
+        cs.restore_x();
+        if (backup_sz < current_sz) {
+            // New columns were added after backup.
+            // move_non_basic_columns_to_bounds snaps non-basic
+            // columns to their bounds and finds a feasible solution.
+            move_non_basic_columns_to_bounds();
+        }
+        else {
+            SASSERT(ax_is_correct());
+            SASSERT(cs.m_r_solver.calc_current_x_is_feasible_include_non_basis());
+        }
+    }
 
     void updt_params(params_ref const& p);
     column_type get_column_type(unsigned j) const { return get_core_solver().m_column_types()[j]; }

src/smt/theory_lra.cpp

@@ -3983,12 +3983,86 @@ public:
         return inf_eps(rational(0), inf_rational(ival.x, ival.y));
     }
 
+    lp::lp_status max_with_lp(theory_var v, lpvar& vi, lp::impq& term_max) {
+        if (!lp().is_feasible() || lp().has_changed_columns())
+            make_feasible();
+        vi = get_lpvar(v);
+        auto st = lp().maximize_term(vi, term_max);
+        if (has_int() && lp().has_inf_int()) {
+            st = lp::lp_status::FEASIBLE;
+            lp().restore_x();
+        }
+        return st;
+    }
+
+    // Returns true if NLA handled the result (blocker and result are set).
+    // Returns false if maximize should fall through to the normal status switch.
+    bool max_with_nl(theory_var v, lp::lp_status& st, unsigned level, expr_ref& blocker, inf_eps& result) {
+        if (!m_nla || (st != lp::lp_status::OPTIMAL && st != lp::lp_status::UNBOUNDED))
+            return false;
+        // Save the LP optimum before NLA check may restore x.
+        auto lp_val = value(v);
+        auto lp_ival = get_ivalue(v);
+        auto nla_st = check_nla(level);
+        TRACE(opt, tout << "check_nla returned " << nla_st 
+              << " lp_ival=" << lp_ival << "\n";
+              if (nla_st == FC_CONTINUE) {
+                  tout << "LP assignment at maximize optimum:\n";
+                  for (unsigned j = 0; j < lp().column_count(); j++) {
+                      if (!lp().get_column_value(j).is_zero())
+                          tout << "  x[" << j << "] = " << lp().get_column_value(j) << "\n";
+                  }
+              });
+        switch (nla_st) {
+        case FC_DONE:
+            // NLA satisfied: keep the optimal assignment, return LP value
+            blocker = mk_gt(v);
+            result = lp_val;
+            st = lp::lp_status::FEASIBLE;
+            return true;
+        case FC_CONTINUE:
+            // NLA found the LP optimum violates nonlinear constraints.
+            // Restore x but return the LP optimum value and blocker
+            // as a bound for the optimizer to validate via check_bound().
+            lp().restore_x();
+            blocker = mk_gt(v, lp_ival);
+            result = lp_val;
+            st = lp::lp_status::FEASIBLE;
+            return true;
+        case FC_GIVEUP:
+            lp().restore_x();
+            st = lp::lp_status::UNBOUNDED;
+            return false;
+        }
+        UNREACHABLE();
+        return false;
+    }
+
+    theory_lra::inf_eps max_result(theory_var v, lpvar vi, lp::lp_status st, expr_ref& blocker, bool& has_shared) {
+        switch (st) {
+        case lp::lp_status::OPTIMAL:
+            init_variable_values();
+            TRACE(arith, display(tout << st << " v" << v << " vi: " << vi << "\n"););
+            blocker = mk_gt(v);
+            return value(v);
+        case lp::lp_status::FEASIBLE:
+            TRACE(arith, display(tout << st << " v" << v << " vi: " << vi << "\n"););
+            blocker = mk_gt(v);
+            return value(v);
+        default:
+            SASSERT(st == lp::lp_status::UNBOUNDED);
+            TRACE(arith, display(tout << st << " v" << v << " vi: " << vi << "\n"););
+            has_shared = false;
+            blocker = m.mk_false();
+            return inf_eps(rational::one(), inf_rational());
+        }
+    }
+
     theory_lra::inf_eps maximize(theory_var v, expr_ref& blocker, bool& has_shared) {
         unsigned level = 2;
         lp::impq term_max;
         lp::lp_status st;
         lpvar vi = 0;
-        unsigned size_of_backup = lp().column_count();
         if (has_int()) {
             lp().backup_x();
         }
@@ -4000,57 +4074,21 @@ public:
             st = lp::lp_status::UNBOUNDED;
         }
         else {
-            if (!lp().is_feasible() || lp().has_changed_columns())
-                make_feasible();
-            
-            vi = get_lpvar(v);
-            
-            st = lp().maximize_term(vi, term_max);
-
-            if (has_int() && lp().has_inf_int()) {
-                st = lp::lp_status::FEASIBLE;
-                if (lp().column_count() == size_of_backup)
-                    lp().restore_x();
-            }
-            if (m_nla && (st == lp::lp_status::OPTIMAL || st == lp::lp_status::UNBOUNDED)) {
-                switch (check_nla(level)) {
-                case FC_DONE:
-                case FC_CONTINUE:
-                    st = lp::lp_status::FEASIBLE;
-                    break;
-                case FC_GIVEUP:
-                    st = lp::lp_status::UNBOUNDED;
-                    break;
-                }                
-                if (lp().column_count() == size_of_backup)
-                    lp().restore_x();
-            }                
-        }
-        switch (st) {
-        case lp::lp_status::OPTIMAL: {
-            init_variable_values();
-            TRACE(arith, display(tout << st << " v" << v << " vi: " << vi << "\n"););
-            auto val = value(v);
-            blocker = mk_gt(v);
-            return val;
-        }
-        case lp::lp_status::FEASIBLE: {
-            auto val = value(v);
-            TRACE(arith, display(tout << st << " v" << v << " vi: " << vi << "\n"););
-            blocker = mk_gt(v);
-            return val;
-        }
-        default:
-            SASSERT(st == lp::lp_status::UNBOUNDED);
-            TRACE(arith, display(tout << st << " v" << v << " vi: " << vi << "\n"););
-            has_shared = false;
-            blocker = m.mk_false();
-            return inf_eps(rational::one(), inf_rational());
+            st = max_with_lp(v, vi, term_max);
+            inf_eps nl_result;
+            if (max_with_nl(v, st, level, blocker, nl_result))
+                return nl_result;
         }
+        return max_result(v, vi, st, blocker, has_shared);
     }
 
     expr_ref mk_gt(theory_var v) {
         lp::impq val = get_ivalue(v);
+        return mk_gt(v, val);
+    }
+
+    // Overload: create blocker from a saved impq value (used when x has been restored)
+    expr_ref mk_gt(theory_var v, lp::impq const& val) {
         expr* obj = get_enode(v)->get_expr();
         rational r = val.x;
         expr_ref e(m);

src/test/api.cpp

@@ -160,9 +160,128 @@ void test_optimize_translate() {
     Z3_del_context(ctx1);
 }
 
+void test_bnh_optimize() {
+    // BNH multi-objective optimization problem using Z3 Optimize C API.
+    // Mimics /tmp/bnh_z3.py: two objectives over a constrained 2D domain.
+    //   f1 = 4*x1^2 + 4*x2^2
+    //   f2 = (x1-5)^2 + (x2-5)^2
+    //   0 <= x1 <= 5, 0 <= x2 <= 3
+    //   C1: (x1-5)^2 + x2^2 <= 25
+    //   C2: (x1-8)^2 + (x2+3)^2 >= 7.7
+
+    Z3_config cfg = Z3_mk_config();
+    Z3_context ctx = Z3_mk_context(cfg);
+    Z3_del_config(cfg);
+
+    Z3_sort real_sort = Z3_mk_real_sort(ctx);
+    Z3_ast x1 = Z3_mk_const(ctx, Z3_mk_string_symbol(ctx, "x1"), real_sort);
+    Z3_ast x2 = Z3_mk_const(ctx, Z3_mk_string_symbol(ctx, "x2"), real_sort);
+
+    auto mk_real = [&](int num, int den = 1) { return Z3_mk_real(ctx, num, den); };
+    auto mk_mul = [&](Z3_ast a, Z3_ast b) { Z3_ast args[] = {a, b}; return Z3_mk_mul(ctx, 2, args); };
+    auto mk_add = [&](Z3_ast a, Z3_ast b) { Z3_ast args[] = {a, b}; return Z3_mk_add(ctx, 2, args); };
+    auto mk_sub = [&](Z3_ast a, Z3_ast b) { Z3_ast args[] = {a, b}; return Z3_mk_sub(ctx, 2, args); };
+    auto mk_sq = [&](Z3_ast a) { return mk_mul(a, a); };
+
+    // f1 = 4*x1^2 + 4*x2^2
+    Z3_ast f1 = mk_add(mk_mul(mk_real(4), mk_sq(x1)), mk_mul(mk_real(4), mk_sq(x2)));
+    // f2 = (x1-5)^2 + (x2-5)^2
+    Z3_ast f2 = mk_add(mk_sq(mk_sub(x1, mk_real(5))), mk_sq(mk_sub(x2, mk_real(5))));
+
+    // Helper: create optimize with BNH constraints and timeout
+    auto mk_bnh_opt = [&]() -> Z3_optimize {
+        Z3_optimize opt = Z3_mk_optimize(ctx);
+        Z3_optimize_inc_ref(ctx, opt);
+        // Set timeout to 5 seconds
+        Z3_params p = Z3_mk_params(ctx);
+        Z3_params_inc_ref(ctx, p);
+        Z3_params_set_uint(ctx, p, Z3_mk_string_symbol(ctx, "timeout"), 5000);
+        Z3_optimize_set_params(ctx, opt, p);
+        Z3_params_dec_ref(ctx, p);
+        // Add BNH constraints
+        Z3_optimize_assert(ctx, opt, Z3_mk_ge(ctx, x1, mk_real(0)));
+        Z3_optimize_assert(ctx, opt, Z3_mk_le(ctx, x1, mk_real(5)));
+        Z3_optimize_assert(ctx, opt, Z3_mk_ge(ctx, x2, mk_real(0)));
+        Z3_optimize_assert(ctx, opt, Z3_mk_le(ctx, x2, mk_real(3)));
+        Z3_optimize_assert(ctx, opt, Z3_mk_le(ctx, mk_add(mk_sq(mk_sub(x1, mk_real(5))), mk_sq(x2)), mk_real(25)));
+        Z3_optimize_assert(ctx, opt, Z3_mk_ge(ctx, mk_add(mk_sq(mk_sub(x1, mk_real(8))), mk_sq(mk_add(x2, mk_real(3)))), mk_real(77, 10)));
+        return opt;
+    };
+
+    auto result_str = [](Z3_lbool r) { return r == Z3_L_TRUE ? "sat" : r == Z3_L_FALSE ? "unsat" : "unknown"; };
+
+    unsigned num_sat = 0;
+
+    // Approach 1: Minimize f1 (Python: opt.minimize(f1))
+    {
+        Z3_optimize opt = mk_bnh_opt();
+        Z3_optimize_minimize(ctx, opt, f1);
+        Z3_lbool result = Z3_optimize_check(ctx, opt, 0, nullptr);
+        std::cout << "BNH min f1: " << result_str(result) << std::endl;
+        if (result == Z3_L_TRUE) {
+            Z3_model m = Z3_optimize_get_model(ctx, opt);
+            Z3_model_inc_ref(ctx, m);
+            Z3_ast val; Z3_model_eval(ctx, m, f1, true, &val);
+            std::cout << "  f1=" << Z3_ast_to_string(ctx, val) << std::endl;
+            Z3_model_dec_ref(ctx, m);
+            num_sat++;
+        }
+        Z3_optimize_dec_ref(ctx, opt);
+    }
+
+    // Approach 2: Minimize f2 (Python: opt2.minimize(f2))
+    {
+        Z3_optimize opt = mk_bnh_opt();
+        Z3_optimize_minimize(ctx, opt, f2);
+        Z3_lbool result = Z3_optimize_check(ctx, opt, 0, nullptr);
+        std::cout << "BNH min f2: " << result_str(result) << std::endl;
+        if (result == Z3_L_TRUE) {
+            Z3_model m = Z3_optimize_get_model(ctx, opt);
+            Z3_model_inc_ref(ctx, m);
+            Z3_ast val; Z3_model_eval(ctx, m, f2, true, &val);
+            std::cout << "  f2=" << Z3_ast_to_string(ctx, val) << std::endl;
+            Z3_model_dec_ref(ctx, m);
+            num_sat++;
+        }
+        Z3_optimize_dec_ref(ctx, opt);
+    }
+
+    // Approach 3: Weighted sum method (Python loop over weights)
+    int weights[][2] = {{1, 4}, {2, 3}, {1, 1}, {3, 2}, {4, 1}};
+    for (auto& w : weights) {
+        Z3_optimize opt = mk_bnh_opt();
+        Z3_ast weighted = mk_add(mk_mul(mk_real(w[0], 100), f1), mk_mul(mk_real(w[1], 100), f2));
+        Z3_optimize_minimize(ctx, opt, weighted);
+        Z3_lbool result = Z3_optimize_check(ctx, opt, 0, nullptr);
+        std::cout << "BNH weighted (w1=" << w[0] << "/5, w2=" << w[1] << "/5): "
+                  << result_str(result) << std::endl;
+        if (result == Z3_L_TRUE) {
+            Z3_model m = Z3_optimize_get_model(ctx, opt);
+            Z3_model_inc_ref(ctx, m);
+            Z3_ast v1, v2;
+            Z3_model_eval(ctx, m, f1, true, &v1);
+            Z3_model_eval(ctx, m, f2, true, &v2);
+            std::cout << "  f1=" << Z3_ast_to_string(ctx, v1)
+                      << " f2=" << Z3_ast_to_string(ctx, v2) << std::endl;
+            Z3_model_dec_ref(ctx, m);
+            num_sat++;
+        }
+        Z3_optimize_dec_ref(ctx, opt);
+    }
+
+    std::cout << "BNH: " << num_sat << "/7 optimizations returned sat" << std::endl;
+    Z3_del_context(ctx);
+    std::cout << "BNH optimization test done" << std::endl;
+}
+
 void tst_api() {
     test_apps();
     test_bvneg();
     test_mk_distinct();
     test_optimize_translate();
+    test_bnh_optimize();
+}
+
+void tst_bnh_opt() {
+    test_bnh_optimize();
 }

src/test/lp/lp.cpp

@@ -564,6 +564,7 @@ void setup_args_parser(argument_parser &parser) {
                                        "test rationals using plus instead of +=");
     parser.add_option_with_help_string("--maximize_term", "test maximize_term()");
     parser.add_option_with_help_string("--patching", "test patching");
+    parser.add_option_with_help_string("--restore_x", "test restore_x");
 }
 
 struct fff {
@@ -1765,6 +1766,124 @@ void test_gomory_cut() {
     
 void test_nla_order_lemma() { nla::test_order_lemma(); }
 
+void test_restore_x() {
+    std::cout << "testing restore_x" << std::endl;
+
+    // Test 1: backup shorter than current (new variables added after backup)
+    {
+        lar_solver solver;
+        lpvar x = solver.add_var(0, false);
+        lpvar y = solver.add_var(1, false);
+        solver.add_var_bound(x, GE, mpq(0));
+        solver.add_var_bound(x, LE, mpq(10));
+        solver.add_var_bound(y, GE, mpq(0));
+        solver.add_var_bound(y, LE, mpq(10));
+
+        vector<std::pair<mpq, lpvar>> coeffs;
+        coeffs.push_back({mpq(1), x});
+        coeffs.push_back({mpq(1), y});
+        unsigned t = solver.add_term(coeffs, 2);
+        solver.add_var_bound(t, GE, mpq(3));
+        solver.add_var_bound(t, LE, mpq(15));
+
+        auto status = solver.solve();
+        SASSERT(status == lp_status::OPTIMAL);
+
+        // Backup the current solution
+        solver.backup_x();
+
+        // Add a new variable with bounds, making the system larger
+        lpvar z = solver.add_var(3, false);
+        solver.add_var_bound(z, GE, mpq(1));
+        solver.add_var_bound(z, LE, mpq(5));
+
+        // restore_x should detect backup < current and call move_non_basic_columns_to_bounds
+        solver.restore_x();
+
+        // The solver should find a feasible solution
+        status = solver.get_status();
+        SASSERT(status == lp_status::OPTIMAL || status == lp_status::FEASIBLE);
+        std::cout << "  test 1 (backup shorter): " << lp_status_to_string(status) << " - PASSED" << std::endl;
+    }
+
+    // Test 2: backup longer than current (columns removed after backup, or pop)
+    {
+        lar_solver solver;
+        lpvar x = solver.add_var(0, false);
+        lpvar y = solver.add_var(1, false);
+        solver.add_var_bound(x, GE, mpq(0));
+        solver.add_var_bound(x, LE, mpq(10));
+        solver.add_var_bound(y, GE, mpq(0));
+        solver.add_var_bound(y, LE, mpq(10));
+
+        vector<std::pair<mpq, lpvar>> coeffs;
+        coeffs.push_back({mpq(1), x});
+        coeffs.push_back({mpq(1), y});
+        unsigned t = solver.add_term(coeffs, 2);
+        solver.add_var_bound(t, GE, mpq(2));
+
+        // Add more variables to make backup larger
+        lpvar z = solver.add_var(3, false);
+        solver.add_var_bound(z, GE, mpq(0));
+        solver.add_var_bound(z, LE, mpq(5));
+
+        auto status = solver.solve();
+        (void)status;
+        SASSERT(status == lp_status::OPTIMAL);
+
+        // Backup with the full system
+        solver.backup_x();
+
+        // restore_x with same-size backup should work fine
+        solver.restore_x();
+        std::cout << "  test 2 (same size backup): PASSED" << std::endl;
+    }
+
+    // Test 3: move_non_basic_columns_to_bounds after solve
+    {
+        lar_solver solver;
+        lpvar x = solver.add_var(0, false);
+        lpvar y = solver.add_var(1, false);
+        solver.add_var_bound(x, GE, mpq(1));
+        solver.add_var_bound(x, LE, mpq(10));
+        solver.add_var_bound(y, GE, mpq(1));
+        solver.add_var_bound(y, LE, mpq(10));
+
+        auto status = solver.solve();
+        SASSERT(status == lp_status::OPTIMAL);
+
+        // Add new constraint: x + y >= 5
+        vector<std::pair<mpq, lpvar>> coeffs;
+        coeffs.push_back({mpq(1), x});
+        coeffs.push_back({mpq(1), y});
+        unsigned t = solver.add_term(coeffs, 2);
+        solver.add_var_bound(t, GE, mpq(5));
+        solver.add_var_bound(t, LE, mpq(15));
+
+        // Add another variable
+        lpvar w = solver.add_var(3, false);
+        solver.add_var_bound(w, GE, mpq(2));
+        solver.add_var_bound(w, LE, mpq(8));
+
+        // Solve expanded system, then move non-basic columns to bounds
+        status = solver.solve();
+        SASSERT(status == lp_status::OPTIMAL);
+        solver.move_non_basic_columns_to_bounds();
+        status = solver.get_status();
+        SASSERT(status == lp_status::OPTIMAL || status == lp_status::FEASIBLE);
+
+        // Verify the model satisfies the constraints
+        std::unordered_map<lpvar, mpq> model;
+        solver.get_model(model);
+        SASSERT(model[x] >= mpq(1) && model[x] <= mpq(10));
+        SASSERT(model[y] >= mpq(1) && model[y] <= mpq(10));
+        SASSERT(model[w] >= mpq(2) && model[w] <= mpq(8));
+        std::cout << "  test 3 (move_non_basic_columns_to_bounds): " << lp_status_to_string(status) << " - PASSED" << std::endl;
+    }
+
+    std::cout << "restore_x tests passed" << std::endl;
+}
+
 void test_lp_local(int argn, char **argv) {
     // initialize_util_module();
     // initialize_numerics_module();
@@ -1792,6 +1911,10 @@ void test_lp_local(int argn, char **argv) {
         test_patching();
         return finalize(0);
     }
+    if (args_parser.option_is_used("--restore_x")) {
+        test_restore_x();
+        return finalize(0);
+    }
     if (args_parser.option_is_used("-nla_cn")) {
 #ifdef Z3DEBUG
         nla::test_cn();

src/test/main.cpp

@@ -175,6 +175,7 @@ int main(int argc, char ** argv) {
     TST(var_subst);
     TST(simple_parser);
     TST(api);
+    TST(bnh_opt);
     TST(api_algebraic);
     TST(api_polynomial);
     TST(api_pb);

src/util/trace_tags.def

@@ -542,6 +542,7 @@ X(Global, isolate_roots_bug, "isolate roots bug")
 X(Global, ite_bug, "ite bug")
 X(Global, lar_solver_feas, "lar solver feas")
 X(Global, lar_solver_inf_heap, "lar solver inf heap")
+X(Global, lar_solver_restore, "lar solver restore")
 X(Global, Lazard, "Lazard")
 X(Global, lcm_bug, "lcm bug")
 X(Global, le_bug, "le bug")