fix: check for non-prime modulus in nmod_poly

Author: oscarbenjamin

src/flint/test/test_all.py

@@ -2101,7 +2101,7 @@ def test_fmpz_mod_poly():
         assert pow(f, 2**60, g) == pow(pow(f, 2**30, g), 2**30, g)
         assert pow(R_gen, 2**60, g) == pow(pow(R_gen, 2**30, g), 2**30, g)
 
-        # Check other typechecks for pow_mod 
+        # Check other typechecks for pow_mod
         assert raises(lambda: pow(f, -2, g), ValueError)
         assert raises(lambda: pow(f, 1, "A"), TypeError)
         assert raises(lambda: pow(f, "A", g), TypeError)
@@ -2536,10 +2536,6 @@ def test_polys():
     for P, S, is_field, characteristic in _all_polys():
 
         composite_characteristic = characteristic != 0 and not characteristic.is_prime()
-        # nmod_poly crashes for many operations with non-prime modulus
-        #     https://github.com/flintlib/python-flint/issues/124
-        # so we can't even test it...
-        nmod_poly_will_crash = type(P(1)) is flint.nmod_poly and composite_characteristic
 
         assert P([S(1)]) == P([1]) == P(P([1])) == P(1)
 
@@ -2684,30 +2680,58 @@ def setbad(obj, i, val):
         assert raises(lambda: P([1, 2, 3]) * None, TypeError)
         assert raises(lambda: None * P([1, 2, 3]), TypeError)
 
-        assert P([1, 2, 1]) // P([1, 1]) == P([1, 1])
-        assert P([1, 2, 1]) % P([1, 1]) == P([0])
-        assert divmod(P([1, 2, 1]), P([1, 1])) == (P([1, 1]), P([0]))
+        if composite_characteristic and type(P(1)) is flint.nmod_poly:
+            # Z/nZ for n not prime
+            #
+            # fmpz_mod_poly and nmod_poly can sometimes compute division with
+            # composite characteristic, but it is not guaranteed to work. For
+            # fmpz_mod_poly, we can detect the failure and raise an exception.
+            # For nmod_poly, we cannot detect the failure and calling e.g.
+            # nmod_poly_divrem would crash the process so for nmod_poly we
+            # raise an exception in all cases if the modulus is not prime.
+            assert raises(lambda: P([1, 2, 1]) // P([1, 1]), DomainError)
+            assert raises(lambda: P([1, 2, 1]) % P([1, 1]), DomainError)
+            assert raises(lambda: divmod(P([1, 2, 1]), P([1, 1])), DomainError)
+
+            assert raises(lambda: 1 // P([1, 1]), DomainError)
+            assert raises(lambda: 1 % P([1, 1]), DomainError)
+            assert raises(lambda: divmod(1, P([1, 1])), DomainError)
+        else:
+            assert P([1, 2, 1]) // P([1, 1]) == P([1, 1])
+            assert P([1, 2, 1]) % P([1, 1]) == P([0])
+            assert divmod(P([1, 2, 1]), P([1, 1])) == (P([1, 1]), P([0]))
+
+            assert 1 // P([1, 1]) == P([0])
+            assert 1 % P([1, 1]) == P([1])
+            assert divmod(1, P([1, 1])) == (P([0]), P([1]))
+
+            assert P([1, 2, 1]) / P([1, 1]) == P([1, 1])
+
+        assert raises(lambda: P([1, 2]) / 0, ZeroDivisionError)
+
+        assert raises(lambda: 1 / P([1, 1]), DomainError)
+        assert raises(lambda: P([1, 2, 1]) / P([1, 2]), DomainError)
 
         if is_field:
             assert P([1, 1]) // 2 == P([S(1)/2, S(1)/2])
             assert P([1, 1]) % 2 == P([0])
+            assert P([2, 2]) / 2 == P([1, 1])
+            assert P([1, 2]) / 2 == P([S(1)/2, 1])
         elif characteristic == 0:
             assert P([1, 1]) // 2 == P([0, 0])
             assert P([1, 1]) % 2 == P([1, 1])
-        elif nmod_poly_will_crash:
-            pass
+            assert P([2, 2]) / 2 == P([1, 1])
+            assert raises(lambda: P([1, 2]) / 2, DomainError)
         else:
             # Z/nZ for n not prime
             if characteristic % 2 == 0:
                 assert raises(lambda: P([1, 1]) // 2, DomainError)
                 assert raises(lambda: P([1, 1]) % 2, DomainError)
+                assert raises(lambda: P([2, 2]) / 2, DomainError)
+                assert raises(lambda: P([1, 2]) / 2, DomainError)
             else:
                 1/0
 
-        assert 1 // P([1, 1]) == P([0])
-        assert 1 % P([1, 1]) == P([1])
-        assert divmod(1, P([1, 1])) == (P([0]), P([1]))
-
         assert raises(lambda: P([1, 2, 1]) // None, TypeError)
         assert raises(lambda: P([1, 2, 1]) % None, TypeError)
         assert raises(lambda: divmod(P([1, 2, 1]), None), TypeError)
@@ -2724,50 +2748,43 @@ def setbad(obj, i, val):
         assert raises(lambda: P([1, 2, 1]) % P([0]), ZeroDivisionError)
         assert raises(lambda: divmod(P([1, 2, 1]), P([0])), ZeroDivisionError)
 
-        # Exact/field scalar division
-        if is_field:
-            assert P([2, 2]) / 2 == P([1, 1])
-            assert P([1, 2]) / 2 == P([S(1)/2, 1])
-        elif characteristic == 0:
-            assert P([2, 2]) / 2 == P([1, 1])
-            assert raises(lambda: P([1, 2]) / 2, DomainError)
-        elif nmod_poly_will_crash:
-            pass
-        else:
-            # Z/nZ for n not prime
-            assert raises(lambda: P([2, 2]) / 2, DomainError)
-            assert raises(lambda: P([1, 2]) / 2, DomainError)
-
-        assert raises(lambda: P([1, 2]) / 0, ZeroDivisionError)
-
-        if not nmod_poly_will_crash:
-            assert P([1, 2, 1]) / P([1, 1]) == P([1, 1])
-            assert raises(lambda: 1 / P([1, 1]), DomainError)
-            assert raises(lambda: P([1, 2, 1]) / P([1, 2]), DomainError)
-
         assert P([1, 1]) ** 0 == P([1])
         assert P([1, 1]) ** 1 == P([1, 1])
         assert P([1, 1]) ** 2 == P([1, 2, 1])
         assert raises(lambda: P([1, 1]) ** -1, ValueError)
         assert raises(lambda: P([1, 1]) ** None, TypeError)
-        
-        # XXX: Not sure what this should do in general:
+
+        # 3-arg pow: (x^2 + 1)**3 mod x-1
+
+        pow3_types = [
+            # flint.fmpq_poly,  XXX
+            flint.nmod_poly,
+            flint.fmpz_mod_poly,
+            flint.fq_default_poly
+        ]
+
         p = P([1, 1])
         mod = P([1, 1])
-        if type(p) not in [flint.fmpz_mod_poly, flint.nmod_poly, flint.fq_default_poly]:
+
+        if type(p) not in pow3_types:
             assert raises(lambda: pow(p, 2, mod), NotImplementedError)
+            assert p * p % mod == 0
+        elif composite_characteristic and type(p) == flint.nmod_poly:
+            # nmod_poly does not support % with composite characteristic
+            assert pow(p, 2, mod) == 0
+            assert raises(lambda: p * p % mod, DomainError)
         else:
+            # Should be for any is_field including fmpq_poly. Works also in
+            # some cases for fmpz_mod_poly with non-prime modulus.
             assert p * p % mod == pow(p, 2, mod)
 
         if not composite_characteristic:
             assert P([1, 2, 1]).gcd(P([1, 1])) == P([1, 1])
-            assert raises(lambda: P([1, 2, 1]).gcd(None), TypeError)
-        elif nmod_poly_will_crash:
-            pass
         else:
             # Z/nZ for n not prime
             assert raises(lambda: P([1, 2, 1]).gcd(P([1, 1])), DomainError)
-            assert raises(lambda: P([1, 2, 1]).gcd(None), TypeError)
+
+        assert raises(lambda: P([1, 2, 1]).gcd(None), TypeError)
 
         if is_field:
             p1 = P([1, 0, 1])
@@ -2778,19 +2795,16 @@ def setbad(obj, i, val):
 
         if not composite_characteristic:
             assert P([1, 2, 1]).factor() == (S(1), [(P([1, 1]), 2)])
-        elif nmod_poly_will_crash:
-            pass
         else:
             assert raises(lambda: P([1, 2, 1]).factor(), DomainError)
 
         if not composite_characteristic:
             assert P([1, 2, 1]).sqrt() == P([1, 1])
-            assert raises(lambda: P([1, 2, 2]).sqrt(), DomainError)
-        elif nmod_poly_will_crash:
-            pass
         else:
             assert raises(lambda: P([1, 2, 1]).sqrt(), DomainError)
 
+        assert raises(lambda: P([1, 2, 2]).sqrt(), DomainError)
+
         if P == flint.fmpq_poly:
             assert raises(lambda: P([1, 2, 1], 3).sqrt(), ValueError)
             assert P([1, 2, 1], 4).sqrt() == P([1, 1], 2)
@@ -3362,13 +3376,6 @@ def factor_sqf(p):
     for P, S, [x, y], is_field, characteristic in _all_polys_mpolys():
 
         if characteristic != 0 and not characteristic.is_prime():
-            # nmod_poly crashes for many operations with non-prime modulus
-            #     https://github.com/flintlib/python-flint/issues/124
-            # so we can't even test it...
-            nmod_poly_will_crash = type(x) is flint.nmod_poly
-            if nmod_poly_will_crash:
-                continue
-
             try:
                 S(4).sqrt() ** 2 == S(4)
             except DomainError:
@@ -4128,7 +4135,7 @@ def test_fq_default():
 
     # p must be prime
     assert raises(lambda: flint.fq_default_ctx(10), ValueError)
-    
+
     # degree must be positive
     assert raises(lambda: flint.fq_default_ctx(11, -1), ValueError)
 

src/flint/types/nmod_poly.pyx

@@ -504,10 +504,14 @@ cdef class nmod_poly(flint_poly):
 
     def _floordiv_(s, t):
         cdef nmod_poly r
+
         if (<nmod_poly>s).val.mod.n != (<nmod_poly>t).val.mod.n:
             raise ValueError("cannot divide nmod_polys with different moduli")
         if nmod_poly_is_zero((<nmod_poly>t).val):
             raise ZeroDivisionError("polynomial division by zero")
+        if not s.ctx._is_prime:
+            raise DomainError("nmod_poly divmod: modulus {self.ctx.mod.n} is not prime")
+
         r = nmod_poly_new_init_preinv(s.ctx)
         nmod_poly_div(r.val, (<nmod_poly>s).val, (<nmod_poly>t).val)
         r.ctx = s.ctx
@@ -527,10 +531,14 @@ cdef class nmod_poly(flint_poly):
 
     def _divmod_(s, t):
         cdef nmod_poly P, Q
+
         if (<nmod_poly>s).val.mod.n != (<nmod_poly>t).val.mod.n:
             raise ValueError("cannot divide nmod_polys with different moduli")
         if nmod_poly_is_zero((<nmod_poly>t).val):
             raise ZeroDivisionError("polynomial division by zero")
+        if not s.ctx._is_prime:
+            raise DomainError("nmod_poly divmod: modulus {self.ctx.mod.n} is not prime")
+
         P = nmod_poly_new_init_preinv(s.ctx)
         Q = nmod_poly_new_init_preinv(s.ctx)
         nmod_poly_divrem(P.val, Q.val, (<nmod_poly>s).val, (<nmod_poly>t).val)
@@ -640,27 +648,53 @@ cdef class nmod_poly(flint_poly):
             >>> (A * B).gcd(B) * 5
             5*x^2 + x + 4
 
+        The modulus must be prime.
         """
         cdef nmod_poly res
+
         other = self.ctx.any_as_nmod_poly(other)
         if other is NotImplemented:
             raise TypeError("cannot convert input to nmod_poly")
         if self.val.mod.n != (<nmod_poly>other).val.mod.n:
             raise ValueError("moduli must be the same")
+        if not self.ctx._is_prime:
+            raise DomainError("nmod_poly gcd: modulus {self.ctx.mod.n} is not prime")
+
         res = nmod_poly_new_init_preinv(self.ctx)
         nmod_poly_gcd(res.val, self.val, (<nmod_poly>other).val)
         res.ctx = self.ctx
         return res
 
     def xgcd(self, other):
+        r"""
+        Computes the extended gcd of self and other: (`G`, `S`, `T`)
+        where `G` is the ``gcd(self, other)`` and `S`, `T` are such that:
+
+        :math:`G = \textrm{self}*S +  \textrm{other}*T`
+
+            >>> f = nmod_poly([143, 19, 37, 138, 102, 127, 95], 163)
+            >>> g = nmod_poly([139, 9, 35, 154, 87, 120, 24], 163)
+            >>> f.xgcd(g)
+            (x^3 + 128*x^2 + 123*x + 91, 17*x^2 + 49*x + 104, 21*x^2 + 5*x + 25)
+
+        The modulus must be prime.
+        """
         cdef nmod_poly res1, res2, res3
+
         other = self.ctx.any_as_nmod_poly(other)
         if other is NotImplemented:
             raise TypeError("cannot convert input to fmpq_poly")
+        if self.val.mod.n != (<nmod_poly>other).val.mod.n:
+            raise ValueError("moduli must be the same")
+        if not self.ctx._is_prime:
+            raise DomainError("nmod_poly xgcd: modulus {self.ctx.mod.n} is not prime")
+
         res1 = nmod_poly_new_init(self.ctx)
         res2 = nmod_poly_new_init(self.ctx)
         res3 = nmod_poly_new_init(self.ctx)
+
         nmod_poly_xgcd(res1.val, res2.val, res3.val, self.val, (<nmod_poly>other).val)
+
         return (res1, res2, res3)
 
     def factor(self, algorithm=None):
@@ -685,11 +719,14 @@ cdef class nmod_poly(flint_poly):
             >>> nmod_poly([3,2,1,2,3], 7).factor(algorithm='cantor-zassenhaus')
             (3, [(x + 4, 1), (x + 2, 1), (x^2 + 4*x + 1, 1)])
 
+        The modulus must be prime.
         """
         if algorithm is None:
             algorithm = 'irreducible'
         elif algorithm not in ('berlekamp', 'cantor-zassenhaus'):
             raise ValueError(f"unknown factorization algorithm: {algorithm}")
+        if not self.ctx._is_prime:
+            raise DomainError(f"nmod_poly factor: modulus {self.ctx.mod.n} is not prime")
         return self._factor(algorithm)
 
     def factor_squarefree(self):
@@ -708,6 +745,8 @@ cdef class nmod_poly(flint_poly):
             (2, [(x, 2), (x + 5, 2), (x + 1, 3)])
 
         """
+        if not self.ctx._is_prime:
+            raise DomainError(f"nmod_poly factor_squarefree: modulus {self.ctx.mod.n} is not prime")
         return self._factor('squarefree')
 
     def _factor(self, factor_type):
@@ -748,12 +787,18 @@ cdef class nmod_poly(flint_poly):
 
     def sqrt(nmod_poly self):
         """Return exact square root or ``None``. """
-        cdef nmod_poly res = nmod_poly_new_init_preinv(self.ctx)
-        if nmod_poly_sqrt(res.val, self.val):
-            return res
-        else:
+        cdef nmod_poly
+
+        if not self.ctx._is_prime:
+            raise DomainError(f"nmod_poly sqrt: modulus {self.ctx.mod.n} is not prime")
+
+        res = nmod_poly_new_init_preinv(self.ctx)
+
+        if not nmod_poly_sqrt(res.val, self.val):
             raise DomainError(f"Cannot compute square root of {self}")
 
+        return res
+
     def deflation(self):
         cdef nmod_poly v
         cdef ulong n