Extend pow_trunc to arbitrary size exponents

Author: remyoudompheng

src/flint/test/test_all.py

@@ -2969,6 +2969,15 @@ def setbad(obj, i, val):
     assert raises(lambda: P([1, 2, 3]).mul_low(None, 3), TypeError) # type: ignore
     assert raises(lambda: P([1, 2, 3]).mul_low(P([4, 5, 6]), None), TypeError) # type: ignore
 
+    p = P([1, 2, 3])
+    assert p.pow_trunc(1234, 3) == P([1, 2468, 3046746])
+    assert raises(lambda: p.pow_trunc(None, 3), TypeError) # type: ignore
+    assert raises(lambda: p.pow_trunc(3, "A"), TypeError) # type: ignore
+    assert raises(lambda: p.pow_trunc(P([4, 5, 6]), 3), TypeError) # type: ignore
+    # Large exponents are allowed
+    assert p.pow_trunc(2**100, 2) == P([1, 2**101])
+    assert p.pow_trunc(6**60, 3) == p.pow_trunc(2**60, 3).pow_trunc(3**60, 3)
+
     # XXX: Not sure what this should do in general:
     p = P([1, 1])
     mod = P([1, 1])

src/flint/types/fmpq_poly.pyx

@@ -526,14 +526,11 @@ cdef class fmpq_poly(flint_poly):
         fmpq_poly_mullow(res.val, self.val, (<fmpq_poly>other).val, n)
         return res
 
-    def pow_trunc(self, slong e, slong n):
+    def pow_trunc(self, e, slong n):
         r"""
         Returns ``self`` raised to the power ``e`` modulo `x^n`:
         :math:`f^e \mod x^n`/
 
-        Note: For exponents larger that 2^63 (which do not fit inside a slong) use the
-        method :meth:`~.pow_mod` with the explicit modulus `x^n`.
-
             >>> f = fmpq_poly([1, 2, 3])
             >>> x = fmpq_poly([0, 1])
             >>> f.pow_trunc(2**20, 4)
@@ -544,7 +541,23 @@ cdef class fmpq_poly(flint_poly):
         if e < 0:
             raise ValueError("Exponent must be non-negative")
 
-        cdef fmpq_poly res
+        cdef slong e_c
+        cdef fmpq_poly res, tmp
+
+        try:
+            e_c = e
+        except OverflowError:
+            # Exponent does not fit slong
+            res = fmpq_poly.__new__(fmpq_poly)
+            tmp = fmpq_poly.__new__(fmpq_poly)
+            ebytes = e.to_bytes((e.bit_length() + 15) // 16 * 2, "big")
+            fmpq_poly_pow_trunc(res.val, self.val, ebytes[0] * 256 + ebytes[1], n)
+            for i in range(2, len(ebytes), 2):
+                fmpq_poly_pow_trunc(res.val, res.val, 1 << 16, n)
+                fmpq_poly_pow_trunc(tmp.val, self.val, ebytes[i] * 256 + ebytes[i+1], n)
+                fmpq_poly_mullow(res.val, res.val, tmp.val, n)
+            return res
+
         res = fmpq_poly.__new__(fmpq_poly)
         fmpq_poly_pow_trunc(res.val, self.val, e, n)
         return res

src/flint/types/fmpz_mod_poly.pyx

@@ -1668,26 +1668,41 @@ cdef class fmpz_mod_poly(flint_poly):
         )
         return res
 
-    def pow_trunc(self, slong e, slong n):
+    def pow_trunc(self, e, slong n):
         r"""
         Returns ``self`` raised to the power ``e`` modulo `x^n`:
         :math:`f^e \mod x^n`/
 
-        Note: For exponents larger that 2^31 (which do not fit inside a ulong) use the
-        method :meth:`~.pow_mod` with the explicit modulus `x^n`.
-
             >>> R = fmpz_mod_poly_ctx(163)
             >>> x = R.gen()
             >>> f = 30*x**6 + 104*x**5 + 76*x**4 + 33*x**3 + 70*x**2 + 44*x + 65
             >>> f.pow_trunc(2**20, 30) == pow(f, 2**20, x**30)
             True
             >>> f.pow_trunc(2**20, 5)
             132*x^4 + 113*x^3 + 36*x^2 + 48*x + 6
+            >>> f.pow_trunc(5**25, 5)
+            147*x^4 + 98*x^3 + 95*x^2 + 33*x + 126
         """
         if e < 0:
             raise ValueError("Exponent must be non-negative")
 
-        cdef fmpz_mod_poly res
+        cdef fmpz_mod_poly res, tmp
+        cdef slong e_c
+
+        try:
+            e_c = e
+        except OverflowError:
+            # Exponent does not fit slong
+            res = self.ctx.new_ctype_poly()
+            tmp = self.ctx.new_ctype_poly()
+            ebytes = e.to_bytes((e.bit_length() + 15) // 16 * 2, "big")
+            fmpz_mod_poly_pow_trunc(res.val, self.val, ebytes[0] * 256 + ebytes[1], n, res.ctx.mod.val)
+            for i in range(2, len(ebytes), 2):
+                fmpz_mod_poly_pow_trunc(res.val, res.val, 1 << 16, n, res.ctx.mod.val)
+                fmpz_mod_poly_pow_trunc(tmp.val, self.val, ebytes[i] * 256 + ebytes[i+1], n, res.ctx.mod.val)
+                fmpz_mod_poly_mullow(res.val, res.val, tmp.val, n, res.ctx.mod.val)
+            return res
+
         res = self.ctx.new_ctype_poly()
         fmpz_mod_poly_pow_trunc(res.val, self.val, e, n, res.ctx.mod.val)
         return res

src/flint/types/fmpz_poly.pyx

@@ -507,14 +507,11 @@ cdef class fmpz_poly(flint_poly):
         fmpz_poly_mullow(res.val, self.val, (<fmpz_poly>other).val, n)
         return res
 
-    def pow_trunc(self, slong e, slong n):
+    def pow_trunc(self, e, slong n):
         r"""
         Returns ``self`` raised to the power ``e`` modulo `x^n`:
         :math:`f^e \mod x^n`/
 
-        Note: For exponents larger that 2^63 (which do not fit inside a slong) use the
-        method :meth:`~.pow_mod` with the explicit modulus `x^n`.
-
             >>> f = fmpz_poly([1, 2, 3])
             >>> x = fmpz_poly([0, 1])
             >>> f.pow_trunc(2**20, 4)
@@ -525,7 +522,23 @@ cdef class fmpz_poly(flint_poly):
         if e < 0:
             raise ValueError("Exponent must be non-negative")
 
-        cdef fmpz_poly res
+        cdef slong e_c
+        cdef fmpz_poly res, tmp
+
+        try:
+            e_c = e
+        except OverflowError:
+            # Exponent does not fit slong
+            res = fmpz_poly.__new__(fmpz_poly)
+            tmp = fmpz_poly.__new__(fmpz_poly)
+            ebytes = e.to_bytes((e.bit_length() + 15) // 16 * 2, "big")
+            fmpz_poly_pow_trunc(res.val, self.val, ebytes[0] * 256 + ebytes[1], n)
+            for i in range(2, len(ebytes), 2):
+                fmpz_poly_pow_trunc(res.val, res.val, 1 << 16, n)
+                fmpz_poly_pow_trunc(tmp.val, self.val, ebytes[i] * 256 + ebytes[i+1], n)
+                fmpz_poly_mullow(res.val, res.val, tmp.val, n)
+            return res
+
         res = fmpz_poly.__new__(fmpz_poly)
         fmpz_poly_pow_trunc(res.val, self.val, e, n)
         return res

src/flint/types/fq_default_poly.pyx

@@ -1190,26 +1190,41 @@ cdef class fq_default_poly(flint_poly):
         )
         return res
 
-    def pow_trunc(self, slong e, slong n):
+    def pow_trunc(self, e, slong n):
         r"""
         Returns ``self`` raised to the power ``e`` modulo `x^n`:
         :math:`f^e \mod x^n`/
 
-        Note: For exponents larger that 2^31 (which do not fit inside a ulong) use the
-        method :meth:`~.pow_mod` with the explicit modulus `x^n`.
-
             >>> R = fq_default_poly_ctx(163)
             >>> x = R.gen()
             >>> f = 30*x**6 + 104*x**5 + 76*x**4 + 33*x**3 + 70*x**2 + 44*x + 65
             >>> f.pow_trunc(2**20, 30) == pow(f, 2**20, x**30)
             True
             >>> f.pow_trunc(2**20, 5)
             132*x^4 + 113*x^3 + 36*x^2 + 48*x + 6
+            >>> f.pow_trunc(5**25, 5)
+            147*x^4 + 98*x^3 + 95*x^2 + 33*x + 126
         """
         if e < 0:
             raise ValueError("Exponent must be non-negative")
 
-        cdef fq_default_poly res
+        cdef slong e_c
+        cdef fq_default_poly res, tmp
+
+        try:
+            e_c = e
+        except OverflowError:
+            # Exponent does not fit slong
+            res = self.ctx.new_ctype_poly()
+            tmp = self.ctx.new_ctype_poly()
+            ebytes = e.to_bytes((e.bit_length() + 15) // 16 * 2, "big")
+            fq_default_poly_pow_trunc(res.val, self.val, ebytes[0] * 256 + ebytes[1], n, res.ctx.field.val)
+            for i in range(2, len(ebytes), 2):
+                fq_default_poly_pow_trunc(res.val, res.val, 1 << 16, n, res.ctx.field.val)
+                fq_default_poly_pow_trunc(tmp.val, self.val, ebytes[i] * 256 + ebytes[i+1], n, res.ctx.field.val)
+                fq_default_poly_mullow(res.val, res.val, tmp.val, n, res.ctx.field.val)
+            return res
+
         res = self.ctx.new_ctype_poly()
         fq_default_poly_pow_trunc(res.val, self.val, e, n, res.ctx.field.val)
         return res

src/flint/types/nmod_poly.pyx

@@ -752,14 +752,11 @@ cdef class nmod_poly(flint_poly):
         nmod_poly_mullow(res.val, self.val, (<nmod_poly>other).val, n)
         return res
 
-    def pow_trunc(self, slong e, slong n):
+    def pow_trunc(self, e, slong n):
         r"""
         Returns ``self`` raised to the power ``e`` modulo `x^n`:
         :math:`f^e \mod x^n`/
 
-        Note: For exponents larger that 2^63 (which do not fit inside a slong) use the
-        method :meth:`~.pow_mod` with the explicit modulus `x^n`.
-
             >>> f = nmod_poly([65, 44, 70, 33, 76, 104, 30], 163)
             >>> x = nmod_poly([0, 1], 163)
             >>> f.pow_trunc(2**20, 30) == pow(f, 2**20, x**30)
@@ -772,7 +769,25 @@ cdef class nmod_poly(flint_poly):
         if e < 0:
             raise ValueError("Exponent must be non-negative")
 
-        cdef nmod_poly res = nmod_poly.__new__(nmod_poly)
+        cdef nmod_poly res, tmp
+        cdef slong e_c
+
+        try:
+            e_c = e
+        except OverflowError:
+            # Exponent does not fit slong
+            res = nmod_poly.__new__(nmod_poly)
+            tmp = nmod_poly.__new__(nmod_poly)
+            nmod_poly_init_preinv(res.val, self.val.mod.n, self.val.mod.ninv)
+            nmod_poly_init_preinv(tmp.val, self.val.mod.n, self.val.mod.ninv)
+            ebytes = e.to_bytes((e.bit_length() + 15) // 16 * 2, "big")
+            nmod_poly_pow_trunc(res.val, self.val, ebytes[0] * 256 + ebytes[1], n)
+            for i in range(2, len(ebytes), 2):
+                nmod_poly_pow_trunc(res.val, res.val, 1 << 16, n)
+                nmod_poly_pow_trunc(tmp.val, self.val, ebytes[i] * 256 + ebytes[i+1], n)
+                nmod_poly_mullow(res.val, res.val, tmp.val, n)
+            return res
+
         res = nmod_poly.__new__(nmod_poly)
         nmod_poly_init_preinv(res.val, self.val.mod.n, self.val.mod.ninv)
         nmod_poly_pow_trunc(res.val, self.val, e, n)