Add missing operations for fmpz and fmpq

oscarbenjamin · oscarbenjamin · commit 08bd191cbaaf · 2023-08-07T18:46:45.000+01:00
diff --git a/src/flint/_flint.pxd b/src/flint/_flint.pxd
@@ -267,6 +267,7 @@ cdef extern from "flint/fmpz.h":
     void fmpz_pow_ui(fmpz_t f,  fmpz_t g, ulong exp)
     void fmpz_powm_ui(fmpz_t f,  fmpz_t g, ulong exp,  fmpz_t m)
     void fmpz_powm(fmpz_t f,  fmpz_t g,  fmpz_t e,  fmpz_t m)
+    int fmpz_pow_fmpz(fmpz_t f, const fmpz_t g, const fmpz_t x)
     int fmpz_sqrtmod(fmpz_t b,  fmpz_t a,  fmpz_t p)
     void fmpz_sqrt(fmpz_t f,  fmpz_t g)
     void fmpz_sqrtrem(fmpz_t f, fmpz_t r,  fmpz_t g)
@@ -310,6 +311,10 @@ cdef extern from "flint/fmpz.h":
     int fmpz_jacobi(const fmpz_t a, const fmpz_t p)
     int fmpz_is_prime(const fmpz_t n)
     int fmpz_is_probabprime(const fmpz_t n)
+    void fmpz_complement(fmpz_t r, const fmpz_t f)
+    void fmpz_and(fmpz_t r, const fmpz_t a, const fmpz_t b)
+    void fmpz_or(fmpz_t r, const fmpz_t a, const fmpz_t b)
+    void fmpz_xor(fmpz_t r, const fmpz_t a, const fmpz_t b)
 
 cdef extern from "flint/fmpz_factor.h":
     ctypedef struct fmpz_factor_struct:
@@ -547,6 +552,7 @@ cdef extern from "flint/fmpq.h":
     void fmpq_div(fmpq_t res, fmpq_t op1, fmpq_t op2)
     void fmpq_div_fmpz(fmpq_t res, fmpq_t op, fmpz_t x)
     int fmpq_mod_fmpz(fmpz_t res, fmpq_t x, fmpz_t mod)
+    int fmpq_pow_fmpz(fmpq_t a, const fmpq_t b, const fmpz_t e)
     int fmpq_reconstruct_fmpz(fmpq_t res, fmpz_t a, fmpz_t m)
     int fmpq_reconstruct_fmpz_2(fmpq_t res, fmpz_t a, fmpz_t m, fmpz_t N, fmpz_t D)
     mp_bitcnt_t fmpq_height_bits(fmpq_t x)
diff --git a/src/flint/fmpq.pyx b/src/flint/fmpq.pyx
@@ -101,6 +101,10 @@ cdef class fmpq(flint_scalar):
     p = property(numer)
     q = property(denom)
 
+    # These are the property names in the numeric tower.
+    numerator = property(numer)
+    denominator = property(denom)
+
     def repr(self):
         if self.q == 1:
             return "fmpq(%s)" % self.p
@@ -122,6 +126,15 @@ cdef class fmpq(flint_scalar):
         else:
             return "%s/%s" % (self.p.str(**kwargs), self.q.str(**kwargs))
 
+    def __int__(self):
+        if self.p >= 0:
+            return int(self.p // self.q)
+        else:
+            return - int((-self.p) // self.q)
+
+    def __floor__(self):
+        return int(self.p // self.q)
+
     def __nonzero__(self):
         return not fmpq_is_zero(self.val)
 
@@ -359,20 +372,28 @@ cdef class fmpq(flint_scalar):
             return max(b1, b2)
 
     def __pow__(self, n, z):
+        cdef fmpz_struct nval[1]
+        cdef int ntype = FMPZ_UNKNOWN
         cdef fmpq v
+        cdef int success
         cdef long e
+
         assert z is None
-        e = n
-        if type(self) is fmpq:
-            v = fmpq.__new__(fmpq)
-            if e >= 0:
-                fmpz_pow_ui(fmpq_numref(v.val), fmpq_numref((<fmpq>self).val), e)
-                fmpz_pow_ui(fmpq_denref(v.val), fmpq_denref((<fmpq>self).val), e)
-            else:
-                if fmpq_is_zero((<fmpq>self).val):
-                    raise ZeroDivisionError
-                fmpz_pow_ui(fmpq_denref(v.val), fmpq_numref((<fmpq>self).val), -e)
-                fmpz_pow_ui(fmpq_numref(v.val), fmpq_denref((<fmpq>self).val), -e)
-            return v
-        return NotImplemented
 
+        ntype = fmpz_set_any_ref(nval, n)
+        if ntype == FMPZ_UNKNOWN:
+            return NotImplemented
+
+        if fmpq_is_zero((<fmpq>self).val) and fmpz_sgn(nval) == -1:
+            if ntype == FMPZ_TMP: fmpz_clear(nval)
+            raise ZeroDivisionError
+
+        v = fmpq.__new__(fmpq)
+        success = fmpq_pow_fmpz(v.val, (<fmpq>self).val, nval)
+
+        if ntype == FMPZ_TMP: fmpz_clear(nval)
+
+        if success:
+            return v
+        else:
+            raise OverflowError("fmpq_pow_fmpz(): exponent too large")
diff --git a/src/flint/fmpz.pyx b/src/flint/fmpz.pyx
@@ -91,6 +91,14 @@ cdef class fmpz(flint_scalar):
                         return
                     raise TypeError("cannot create fmpz from type %s" % type(val))
 
+    @property
+    def numerator(self):
+        return self
+
+    @property
+    def denominator(self):
+        return fmpz(1)
+
     # XXX: improve!
     def __int__(self):
         return fmpz_get_intlong(self.val)
@@ -101,6 +109,9 @@ cdef class fmpz(flint_scalar):
     def __index__(self):
         return fmpz_get_intlong(self.val)
 
+    def __float__(self):
+        return float(fmpz_get_intlong(self.val))
+
     def __richcmp__(s, t, int op):
         cdef bint res = 0
         cdef long tl
@@ -334,28 +345,171 @@ cdef class fmpz(flint_scalar):
         return u
 
     def __pow__(s, t, m):
-        cdef ulong exp
-        u = NotImplemented
-        if m is not None:
-            raise NotImplementedError("modular exponentiation")
-        c = t
-        u = fmpz.__new__(fmpz)
-        fmpz_pow_ui((<fmpz>u).val, (<fmpz>s).val, c)
-        return u
-
-    def __rpow__(s, t, m):
         cdef fmpz_struct tval[1]
-        cdef int stype = FMPZ_UNKNOWN
-        cdef ulong exp
+        cdef fmpz_struct mval[1]
+        cdef int ttype = FMPZ_UNKNOWN
+        cdef int mtype = FMPZ_UNKNOWN
+        cdef int success
         u = NotImplemented
-        if m is not None:
-            raise NotImplementedError("modular exponentiation")
         ttype = fmpz_set_any_ref(tval, t)
-        if ttype != FMPZ_UNKNOWN:
+        if ttype == FMPZ_UNKNOWN:
+            return NotImplemented
+
+        if m is None:
+            # fmpz_pow_fmpz throws if x is negative
+            if fmpz_sgn(tval) == -1:
+                if ttype == FMPZ_TMP: fmpz_clear(tval)
+                raise ValueError("negative exponent")
+
             u = fmpz.__new__(fmpz)
-            s_ulong = fmpz_get_ui(s.val)
-            fmpz_pow_ui((<fmpz>u).val, tval, s_ulong)
+            success = fmpz_pow_fmpz((<fmpz>u).val, (<fmpz>s).val, tval)
+
+            if not success:
+                if ttype == FMPZ_TMP: fmpz_clear(tval)
+                raise ValueError("fmpz_pow_fmpz: exponent too large")
+        else:
+            # Modular exponentiation
+            mtype = fmpz_set_any_ref(mval, m)
+            if mtype != FMPZ_UNKNOWN:
+
+                if fmpz_is_zero(mval):
+                    if ttype == FMPZ_TMP: fmpz_clear(tval)
+                    if mtype == FMPZ_TMP: fmpz_clear(mval)
+                    raise ValueError("pow(): modulus cannot be zero")
+
+                # The Flint docs say that fmpz_powm will throw if m is zero
+                # but it also throws if m is negative. Python generally allows
+                # e.g. pow(2, 2, -3) == (2^2) % (-3) == -2. We could implement
+                # that here as well but it is not clear how useful it is.
+                if fmpz_sgn(mval) == -1:
+                    if ttype == FMPZ_TMP: fmpz_clear(tval)
+                    if mtype == FMPZ_TMP: fmpz_clear(mval)
+                    raise ValueError("pow(): negative modulua not supported")
+
+                u = fmpz.__new__(fmpz)
+                fmpz_powm((<fmpz>u).val, (<fmpz>s).val, tval, mval)
+
         if ttype == FMPZ_TMP: fmpz_clear(tval)
+        if mtype == FMPZ_TMP: fmpz_clear(mval)
+        return u
+
+    def __rpow__(s, t, m):
+        t = any_as_fmpz(t)
+        if t is NotImplemented:
+            return t
+        return t.__pow__(s, m)
+
+    def __lshift__(self, other):
+        if typecheck(other, fmpz):
+            other = int(other)
+        if typecheck(other, int):
+            u = fmpz.__new__(fmpz)
+            fmpz_mul_2exp((<fmpz>u).val, self.val, other)
+            return u
+        else:
+            return NotImplemented
+
+    def __rlshift__(self, other):
+        if typecheck(other, int):
+            u = fmpz.__new__(fmpz)
+            fmpz_mul_2exp((<fmpz>u).val, fmpz(other).val, int(self))
+            return u
+        else:
+            return NotImplemented
+
+    def __rshift__(self, other):
+        if typecheck(other, fmpz):
+            other = int(other)
+        if typecheck(other, int):
+            u = fmpz.__new__(fmpz)
+            fmpz_fdiv_q_2exp((<fmpz>u).val, self.val, other)
+            return u
+        else:
+            return NotImplemented
+
+    def __rrshift__(self, other):
+        if typecheck(other, int):
+            u = fmpz.__new__(fmpz)
+            fmpz_fdiv_q_2exp((<fmpz>u).val, fmpz(other).val, int(self))
+            return u
+        else:
+            return NotImplemented
+
+    def __and__(self, other):
+        cdef fmpz_struct tval[1]
+        cdef int ttype = FMPZ_UNKNOWN
+        ttype = fmpz_set_any_ref(tval, other)
+        if ttype == FMPZ_UNKNOWN:
+            return NotImplemented
+        u = fmpz.__new__(fmpz)
+        fmpz_and((<fmpz>u).val, self.val, tval)
+        if ttype == FMPZ_TMP:
+            fmpz_clear(tval)
+        return u
+
+    def __rand__(self, other):
+        cdef fmpz_struct tval[1]
+        cdef int ttype = FMPZ_UNKNOWN
+        ttype = fmpz_set_any_ref(tval, other)
+        if ttype == FMPZ_UNKNOWN:
+            return NotImplemented
+        u = fmpz.__new__(fmpz)
+        fmpz_and((<fmpz>u).val, tval, self.val)
+        if ttype == FMPZ_TMP:
+            fmpz_clear(tval)
+        return u
+
+    def __or__(self, other):
+        cdef fmpz_struct tval[1]
+        cdef int ttype = FMPZ_UNKNOWN
+        ttype = fmpz_set_any_ref(tval, other)
+        if ttype == FMPZ_UNKNOWN:
+            return NotImplemented
+        u = fmpz.__new__(fmpz)
+        fmpz_or((<fmpz>u).val, self.val, tval)
+        if ttype == FMPZ_TMP:
+            fmpz_clear(tval)
+        return u
+
+    def __ror__(self, other):
+        cdef fmpz_struct tval[1]
+        cdef int ttype = FMPZ_UNKNOWN
+        ttype = fmpz_set_any_ref(tval, other)
+        if ttype == FMPZ_UNKNOWN:
+            return NotImplemented
+        u = fmpz.__new__(fmpz)
+        fmpz_or((<fmpz>u).val, tval, self.val)
+        if ttype == FMPZ_TMP:
+            fmpz_clear(tval)
+        return u
+
+    def __xor__(self, other):
+        cdef fmpz_struct tval[1]
+        cdef int ttype = FMPZ_UNKNOWN
+        ttype = fmpz_set_any_ref(tval, other)
+        if ttype == FMPZ_UNKNOWN:
+            return NotImplemented
+        u = fmpz.__new__(fmpz)
+        fmpz_xor((<fmpz>u).val, self.val, tval)
+        if ttype == FMPZ_TMP:
+            fmpz_clear(tval)
+        return u
+
+    def __rxor__(self, other):
+        cdef fmpz_struct tval[1]
+        cdef int ttype = FMPZ_UNKNOWN
+        ttype = fmpz_set_any_ref(tval, other)
+        if ttype == FMPZ_UNKNOWN:
+            return NotImplemented
+        u = fmpz.__new__(fmpz)
+        fmpz_xor((<fmpz>u).val, tval, self.val)
+        if ttype == FMPZ_TMP:
+            fmpz_clear(tval)
+        return u
+
+    def __invert__(self):
+        u = fmpz.__new__(fmpz)
+        fmpz_complement((<fmpz>u).val, self.val)
         return u
 
     def gcd(self, other):
diff --git a/test/test.py b/test/test.py
@@ -101,6 +101,23 @@ def test_fmpz():
                     assert (ltype(s) >= rtype(t)) == (s >= t)
                     if 0 <= t < 10:
                         assert (ltype(s) ** rtype(t)) == (s ** t)
+
+    pow_mod_examples = [
+        (2, 2, 3, 1),
+        (2, -1, 5, 3),
+        (2, 0, 5, 1),
+    ]
+    for a, b, c, ab_mod_c in pow_mod_examples:
+        assert pow(a, b, c) == ab_mod_c
+        assert pow(flint.fmpz(a), b, c) == ab_mod_c
+        assert pow(a, flint.fmpz(b), c) == ab_mod_c
+        assert pow(flint.fmpz(a), flint.fmpz(b), c) == ab_mod_c
+        assert pow(flint.fmpz(a), flint.fmpz(b), flint.fmpz(c)) == ab_mod_c
+
+    assert raises(lambda: pow(flint.fmpz(2), 2, 0), ValueError)
+    # XXX: Handle negative modulus like int?
+    assert raises(lambda: pow(flint.fmpz(2), 2, -1), ValueError)
+
     assert flint.fmpz(2) != []
     assert +flint.fmpz(0) == 0
     assert +flint.fmpz(1) == 1
@@ -139,7 +156,6 @@ def test_fmpz():
     big = flint.fmpz(bigstr)
     assert big.str() == bigstr
     assert big.str(condense=10) == '1111111111{...80 digits...}1111111111'
-    assert raises(lambda: pow(flint.fmpz(2), 2, 3), NotImplementedError)
 
 def test_fmpz_factor():
     assert flint.fmpz(6).gcd(flint.fmpz(9)) == 3
