From 47d6663c679314776034ba14219a885cfa9a0781 Mon Sep 17 00:00:00 2001
From: Jakob Rath <git@jakobrath.eu>
Date: Mon, 13 Mar 2023 09:29:21 +0100
Subject: [PATCH] support other ops

---
 src/math/polysat/polysat_ast.cpp | 47 +++++++++++++++++++++++++-------
 src/math/polysat/polysat_ast.h   | 14 +++++++---
 2 files changed, 47 insertions(+), 14 deletions(-)

diff --git a/src/math/polysat/polysat_ast.cpp b/src/math/polysat/polysat_ast.cpp
index 0ab84d2f9..ad808a0e3 100644
--- a/src/math/polysat/polysat_ast.cpp
+++ b/src/math/polysat/polysat_ast.cpp
@@ -34,7 +34,11 @@ namespace polysat {
         // just keep all expressions alive
         ast_ref_vector storage;
 
-        polysat_ast_d(): m(), decls(m), formulas(m), storage(m)
+        polysat_ast_d()
+            : m()
+            , decls(m)
+            , formulas(m)
+            , storage(m)
         {
             reg_decl_plugins(m);
             bv = alloc(bv_util, m);
@@ -153,24 +157,47 @@ namespace polysat {
         expr* v_inv_max = bv().mk_bv_sub(bv().mk_bv_shl(one, bv().mk_bv_sub(mk_val(rational(N), N), parity)), one);
         expr* definition2 = bv().mk_ule(v_inv, v_inv_max);
         // conditional on b
-        expr* b = mk_fresh_bool_const("parity_enable");
+        expr* b = mk_fresh_bool_const("op_inv");
         add(m().mk_iff(b, definition1));
         add(m().mk_iff(b, definition2));
         return b;
     }
 
-    expr* polysat_ast::mk_op(op_constraint const& op, bool sign) {
-        expr* e = nullptr;
+    template <typename Fn>
+    expr* polysat_ast::mk_bin(char const* name, pdd const& p, pdd const& q, pdd const& r, Fn mk_bin_expr) {
+        unsigned const N = p.power_of_2();
+        // r = p OP q
+        expr* definition = m().mk_eq(mk_poly(r), mk_bin_expr(mk_poly(p), mk_poly(q)));
+        // b <=> definition
+        expr* b = mk_fresh_bool_const(name);
+        add(m().mk_iff(b, definition));
+        return b;
+    }
+
+    expr* polysat_ast::mk_op(op_constraint const& op) {
+        using code = op_constraint::code;
         switch (op.get_op()) {
-        case op_constraint::code::inv_op:
-            e = mk_inv(op.p(), op.r());
-            break;
+        case code::lshr_op:
+            return mk_bin("op_lshr", op.p(), op.q(), op.r(), [&](expr* p, expr* q) { return bv().mk_bv_lshr(p, q); });
+        case code::ashr_op:
+            return mk_bin("op_ashr", op.p(), op.q(), op.r(), [&](expr* p, expr* q) { return bv().mk_bv_ashr(p, q); });
+        case code::shl_op:
+            return mk_bin("op_shl", op.p(), op.q(), op.r(), [&](expr* p, expr* q) { return bv().mk_bv_shl(p, q); });
+        case code::and_op:
+            return mk_bin("op_and", op.p(), op.q(), op.r(), [&](expr* p, expr* q) { return bv().mk_bv_and(p, q); });
+        case code::inv_op:
+            return mk_inv(op.p(), op.r());
         default:
             return nullptr;
         }
-        if (sign)
-            e = mk_not(e);
-        return e;
+    }
+
+    expr* polysat_ast::mk_op(op_constraint const& op, bool sign) {
+        expr* e = mk_op(op);
+        if (e && sign)
+            return mk_not(e);
+        else
+            return e;
     }
 
     expr* polysat_ast::mk_not(expr* e) {
diff --git a/src/math/polysat/polysat_ast.h b/src/math/polysat/polysat_ast.h
index 2edf57671..a7352e7ed 100644
--- a/src/math/polysat/polysat_ast.h
+++ b/src/math/polysat/polysat_ast.h
@@ -35,9 +35,19 @@ namespace polysat {
 
         expr* mk_fresh_bool_const(char const* prefix);
 
+        expr* mk_ule(pdd const& p, pdd const& q);
+
         // expression for "parity of p"
         expr* mk_parity(pdd const& p);
 
+        // p_inv = pseudo-inverse(p)
+        expr* mk_inv(pdd const& p, pdd const& p_inv);
+
+        expr* mk_op(op_constraint const& op);
+
+        template <typename Fn>
+        expr* mk_bin(char const* name, pdd const& p, pdd const& q, pdd const& r, Fn mk_bin_expr);
+
     public:
         polysat_ast(solver&);
         ~polysat_ast();
@@ -52,7 +62,6 @@ namespace polysat {
         expr* mk_poly(pdd const& p);
 
         // p <= q
-        expr* mk_ule(pdd const& p, pdd const& q);
         expr* mk_ule(pdd const& p, pdd const& q, bool sign);
         // p >= q
         expr* mk_uge(pdd const& p, pdd const& q, bool sign) { return mk_ule(q, p, sign); }
@@ -64,9 +73,6 @@ namespace polysat {
         // ovfl*(p, q)
         expr* mk_umul_ovfl(pdd const& p, pdd const& q, bool sign);
 
-        // p_inv = pseudo-inverse(p)
-        expr* mk_inv(pdd const& p, pdd const& p_inv);
-
         expr* mk_op(op_constraint const& op, bool sign);
 
         expr* mk_not(expr* e);