From e580c384b89c89f6536cc3e2e6a72471b0139bbe Mon Sep 17 00:00:00 2001
From: Nikolaj Bjorner <nbjorner@microsoft.com>
Date: Sat, 9 Dec 2023 12:46:06 -0800
Subject: [PATCH] import updates to rational from polysat

Signed-off-by: Nikolaj Bjorner <nbjorner@microsoft.com>
---
 src/util/mpq.cpp    |  6 ++++++
 src/util/mpq.h      |  8 ++++++++
 src/util/mpz.cpp    | 13 +++++++++++++
 src/util/mpz.h      |  7 +++++++
 src/util/rational.h | 27 ++++++++++++++++++++++++++-
 5 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/src/util/mpq.cpp b/src/util/mpq.cpp
index 324750cfa..f90843e36 100644
--- a/src/util/mpq.cpp
+++ b/src/util/mpq.cpp
@@ -316,6 +316,12 @@ unsigned mpq_manager<SYNCH>::prev_power_of_two(mpq const & a) {
     return prev_power_of_two(_tmp);
 }
 
+template<bool SYNCH>
+unsigned mpq_manager<SYNCH>::next_power_of_two(mpq const & a) {
+    _scoped_numeral<mpz_manager<SYNCH> > _tmp(*this);
+    ceil(a, _tmp);
+    return next_power_of_two(_tmp);
+}
 
 template<bool SYNCH>
 template<bool SUB>
diff --git a/src/util/mpq.h b/src/util/mpq.h
index e254ade69..1bdf8f31b 100644
--- a/src/util/mpq.h
+++ b/src/util/mpq.h
@@ -848,6 +848,14 @@ public:
     unsigned prev_power_of_two(mpz const & a) { return mpz_manager<SYNCH>::prev_power_of_two(a); }
     unsigned prev_power_of_two(mpq const & a);
 
+    /**
+       \brief Return the smallest k s.t. a <= 2^k.
+
+       \remark Return 0 if a is not positive.
+    */
+    unsigned next_power_of_two(mpz const & a) { return mpz_manager<SYNCH>::next_power_of_two(a); }
+    unsigned next_power_of_two(mpq const & a);
+
     bool is_int_perfect_square(mpq const & a, mpq & r) {
         SASSERT(is_int(a));
         reset_denominator(r);
diff --git a/src/util/mpz.cpp b/src/util/mpz.cpp
index c3ba30161..296b4426e 100644
--- a/src/util/mpz.cpp
+++ b/src/util/mpz.cpp
@@ -2288,6 +2288,19 @@ unsigned mpz_manager<SYNCH>::bitsize(mpz const & a) {
         return mlog2(a) + 1;
 }
 
+template<bool SYNCH>
+unsigned mpz_manager<SYNCH>::next_power_of_two(mpz const & a) {
+    if (is_nonpos(a))
+        return 0;
+    if (is_one(a))
+        return 0;
+    unsigned shift;
+    if (is_power_of_two(a, shift))
+        return shift;
+    else
+        return log2(a) + 1;
+}
+
 template<bool SYNCH>
 bool mpz_manager<SYNCH>::is_perfect_square(mpz const & a, mpz & root) {
     if (is_neg(a))
diff --git a/src/util/mpz.h b/src/util/mpz.h
index a1bb19395..bb1141ea7 100644
--- a/src/util/mpz.h
+++ b/src/util/mpz.h
@@ -692,6 +692,13 @@ public:
        \remark Return 0 if a is not positive.
     */
     unsigned prev_power_of_two(mpz const & a) { return log2(a); }
+
+    /**
+       \brief Return the smallest k s.t. a <= 2^k.
+
+       \remark Return 0 if a is not positive.
+    */
+    unsigned next_power_of_two(mpz const & a);
     
     /**
        \brief Return true if a^{1/n} is an integer, and store the result in a.
diff --git a/src/util/rational.h b/src/util/rational.h
index f47fddefe..88a0552ba 100644
--- a/src/util/rational.h
+++ b/src/util/rational.h
@@ -55,7 +55,7 @@ public:
     explicit rational(double  z) { UNREACHABLE(); }
     
     explicit rational(char const * v) { m().set(m_val, v); }
-
+    
     explicit rational(unsigned const * v, unsigned sz) { m().set(m_val, sz, v); }
 
     struct i64 {};
@@ -489,6 +489,18 @@ public:
         return get_num_digits(rational(10));
     }
 
+    /**
+     * \brief Return the biggest k s.t. 2^k <= a.
+     * \remark Return 0 if a is not positive.
+     */
+    unsigned prev_power_of_two() const { return m().prev_power_of_two(m_val); }
+
+    /**
+     * \brief Return the smallest k s.t. a <= 2^k.
+     * \remark Return 0 if a is not positive.
+     */
+    unsigned next_power_of_two() const { return m().next_power_of_two(m_val); }
+
     bool get_bit(unsigned index) const {
         return m().get_bit(m_val, index);
     }
@@ -501,6 +513,15 @@ public:
         return k;
     }
 
+    /** Number of trailing zeros in an N-bit representation */
+    unsigned parity(unsigned num_bits) const {
+        SASSERT(!is_neg());
+        SASSERT(*this < rational::power_of_two(num_bits));
+        if (is_zero())
+            return num_bits;
+        return trailing_zeros();
+    }
+
     static bool limit_denominator(rational &num, rational const& limit);
 };
 
@@ -649,3 +670,7 @@ inline rational gcd(rational const & r1, rational const & r2, rational & a, rati
   rational::m().gcd(r1.m_val, r2.m_val, a.m_val, b.m_val, result.m_val);
   return result;
 }
+
+inline void swap(rational& r1, rational& r2) {
+    r1.swap(r2);
+}