numbertheory.py: generalise SqrtModP to do other roots.

I'm about to want to solve quartics mod a prime, which means I'll need
to be able to take cube roots mod p as well as square roots.

This commit introduces a more general class which can take rth roots
for any prime r, and moreover, it can do it in a general cyclic group.
(You have to tell it the group's order and give it some primitives for
doing arithmetic, plus a way of iterating over the group elements that
it can use to look for a non-rth-power and roots of unity.)

That system makes it nicely easy to test, because you can give it a
cyclic group represented as the integers under _addition_, and then
you obviously know what all the right answers are. So I've also added
a unit test system checking that.

test/eccref.py

@@ -111,9 +111,9 @@ class WeierstrassCurve(CurveBase):
 
     def cpoint(self, x, yparity=0):
         if not hasattr(self, 'sqrtmodp'):
-            self.sqrtmodp = SqrtModP(self.p)
+            self.sqrtmodp = RootModP(2, self.p)
         rhs = x**3 + self.a.n * x + self.b.n
-        y = self.sqrtmodp.sqrt(rhs)
+        y = self.sqrtmodp.root(rhs)
         if (y - yparity) % 2:
             y = -y
         return self.point(x, y)
@@ -157,9 +157,9 @@ class MontgomeryCurve(CurveBase):
 
     def cpoint(self, x, yparity=0):
         if not hasattr(self, 'sqrtmodp'):
-            self.sqrtmodp = SqrtModP(self.p)
+            self.sqrtmodp = RootModP(2, self.p)
         rhs = (x**3 + self.a.n * x**2 + x) / self.b
-        y = self.sqrtmodp.sqrt(int(rhs))
+        y = self.sqrtmodp.root(int(rhs))
         if (y - yparity) % 2:
             y = -y
         return self.point(x, y)
@@ -198,11 +198,11 @@ class TwistedEdwardsCurve(CurveBase):
 
     def cpoint(self, y, xparity=0):
         if not hasattr(self, 'sqrtmodp'):
-            self.sqrtmodp = SqrtModP(self.p)
+            self.sqrtmodp = RootModP(self.p)
         y = ModP(self.p, y)
         y2 = y**2
         radicand = (y2 - 1) / (self.d * y2 - self.a)
-        x = self.sqrtmodp.sqrt(radicand.n)
+        x = self.sqrtmodp.root(radicand.n)
         if (x - xparity) % 2:
             x = -x
         return self.point(x, y)

test/numbertheory.py

@@ -1,5 +1,6 @@
 import numbers
 import itertools
+import unittest
 
 def invert(a, b):
     "Multiplicative inverse of a mod b. a,b must be coprime."
@@ -36,57 +37,148 @@ def jacobi(n,m):
             acc *= -1
         n, m = m, n
 
-class SqrtModP(object):
+class CyclicGroupRootFinder(object):
-    """Class for finding square roots of numbers mod p.
+    """Class for finding rth roots in a cyclic group. r must be prime."""
 
-    p must be an odd prime (but its primality is not checked)."""
+    # Basic strategy:
+    #
+    # We write |G| = r^k u, with u coprime to r. This gives us a
+    # nested sequence of subgroups G = G_0 > G_1 > ... > G_k, each
+    # with index 3 in its predecessor. G_0 is the whole group, and the
+    # innermost G_k has order u.
+    #
+    # Within G_k, you can take an rth root by raising an element to
+    # the power of (r^{-1} mod u). If k=0 (so G = G_0 = G_k) then
+    # that's all that's needed: every element has a unique rth root.
+    # But if k>0, then things go differently.
+    #
+    # Define the 'rank' of an element g as the highest i such that
+    # g \in G_i. Elements of rank 0 are the non-rth-powers: they don't
+    # even _have_ an rth root. Elements of rank k are the easy ones to
+    # take rth roots of, as above.
+    #
+    # In between, you can follow an inductive process, as long as you
+    # know one element z of index 0. Suppose we're trying to take the
+    # rth root of some g with index i. Repeatedly multiply g by
+    # z^{r^i} until its index increases; then take the root of that
+    # (recursively), and divide off z^{r^{i-1}} once you're done.
 
-    def __init__(self, p):
+    def __init__(self, r, order):
-        p = abs(p)
+        self.order = order # order of G
-        assert p & 1
+        self.r = r
-        self.p = p
+        self.k = next(k for k in itertools.count()
+                      if self.order % (r**(k+1)) != 0)
+        self.u = self.order // (r**self.k)
+        self.z = next(z for z in self.iter_elements()
+                      if self.index(z) == 0)
+        self.zinv = self.inverse(self.z)
+        self.root_power = invert(self.r, self.u) if self.u > 1 else 0
 
-        # Decompose p as 2^e k + 1 for odd k.
+        self.roots_of_unity = {self.identity()}
-        self.k = p-1
+        if self.k > 0:
-        self.e = 0
+            exponent = self.order // self.r
-        while not (self.k & 1):
+            for z in self.iter_elements():
-            self.k >>= 1
+                root_of_unity = self.pow(z, exponent)
-            self.e += 1
+                if root_of_unity not in self.roots_of_unity:
+                    self.roots_of_unity.add(root_of_unity)
+                    if len(self.roots_of_unity) == r:
+                        break
 
-        # Find a non-square mod p.
+    def index(self, g):
-        for self.z in itertools.count(1):
+        h = self.pow(g, self.u)
-            if jacobi(self.z, self.p) == -1:
+        for i in range(self.k+1):
-                break
+            if h == self.identity():
-        self.zinv = ModP(self.p, self.z).invert()
+                return self.k - i
+            h = self.pow(h, self.r)
+        assert False, ("Not a cyclic group! Raising {} to u r^k should give e."
+                       .format(g))
 
-    def sqrt_recurse(self, a):
+    def all_roots(self, g):
-        ak = pow(a, self.k, self.p)
+        try:
-        for i in range(self.e, -1, -1):
+            r = self.root(g)
-            if ak == 1:
+        except ValueError:
-                break
+            return []
-            ak = ak*ak % self.p
+        return {r * rou for rou in self.roots_of_unity}
-        assert i > 0
-        if i == self.e:
-            return pow(a, (self.k+1) // 2, self.p)
-        r_prime = self.sqrt_recurse(a * pow(self.z, 2**i, self.p))
-        return r_prime * pow(self.zinv, 2**(i-1), self.p) % self.p
 
-    def sqrt(self, a):
+    def root(self, g):
-        j = jacobi(a, self.p)
+        i = self.index(g)
-        if j == 0:
+        if i == 0 and self.k > 0:
-            return 0
+            raise ValueError("{} has no {}th root".format(g, self.r))
-        if j < 0:
+        out = self.root_recurse(g, i)
-            raise ValueError("{} has no square root mod {}".format(a, self.p))
+        assert self.pow(out, self.r) == g
-        a %= self.p
+        return out
-        r = self.sqrt_recurse(a)
-        assert r*r % self.p == a
-        # Normalise to the smaller (or 'positive') one of the two roots.
-        return min(r, self.p - r)
 
-    def __str__(self):
+    def root_recurse(self, g, i):
-        return "{}({})".format(type(self).__name__, self.p)
+        if i == self.k:
-    def __repr__(self):
+            return self.pow(g, self.root_power)
-        return self.__str__()
+        z_in = self.pow(self.z, self.r**i)
+        z_out = self.pow(self.zinv, self.r**(i-1))
+        adjust = self.identity()
+        while True:
+            g = self.mul(g, z_in)
+            adjust = self.mul(adjust, z_out)
+            i2 = self.index(g)
+            if i2 > i:
+                return self.mul(self.root_recurse(g, i2), adjust)
+
+class AdditiveGroupRootFinder(CyclicGroupRootFinder):
+    """Trivial test subclass for CyclicGroupRootFinder.
+
+    Represents a cyclic group of any order additively, as the integers
+    mod n under addition. This makes root-finding trivial without
+    having to use the complicated algorithm above, and therefore it's
+    a good way to test the complicated algorithm under conditions
+    where the right answers are obvious."""
+
+    def __init__(self, r, order):
+        super().__init__(r, order)
+
+    def mul(self, x, y):
+        return (x + y) % self.order
+    def pow(self, x, n):
+        return (x * n) % self.order
+    def inverse(self, x):
+        return (-x) % self.order
+    def identity(self):
+        return 0
+    def iter_elements(self):
+        return range(self.order)
+
+class TestCyclicGroupRootFinder(unittest.TestCase):
+    def testRootFinding(self):
+        for order in 10, 11, 12, 18:
+            grf = AdditiveGroupRootFinder(3, order)
+            for i in range(order):
+                try:
+                    r = grf.root(i)
+                except ValueError:
+                    r = None
+
+                if order % 3 == 0 and i % 3 != 0:
+                    self.assertEqual(r, None)
+                else:
+                    self.assertEqual(r*3 % order, i)
+
+class RootModP(CyclicGroupRootFinder):
+    """The live class that can take rth roots mod a prime."""
+
+    def __init__(self, r, p):
+        self.modulus = p
+        super().__init__(r, p-1)
+
+    def mul(self, x, y):
+        return (x * y) % self.modulus
+    def pow(self, x, n):
+        return pow(x, n, self.modulus)
+    def inverse(self, x):
+        return invert(x, self.modulus)
+    def identity(self):
+        return 1
+    def iter_elements(self):
+        return range(1, self.modulus)
+
+    def root(self, g):
+        return 0 if g == 0 else super().root(g)
 
 class ModP(object):
     """Class that represents integers mod p as a field.
@@ -179,3 +271,9 @@ class ModP(object):
         return "{}(0x{:x},0x{:x})".format(type(self).__name__, self.p, self.n)
     def __hash__(self):
         return hash((type(self).__name__, self.p, self.n))
+
+if __name__ == "__main__":
+    import sys
+    if sys.argv[1:] == ["--test"]:
+        sys.argv[1:2] = []
+        unittest.main()