implement Tonelli Shanks

TonelliShanks.hs

@@ -0,0 +1,51 @@
+module TonelliShanks (tonelliShanks) where
+
+import ModularArithmeticUtils (modExp, modMul, legendre)
+
+-- Factor p-1 as q * 2^s, where q is odd
+factorOutTwos :: Integer -> (Integer, Integer)
+factorOutTwos n = go n 0
+  where
+    go x s
+      | even x    = go (x `div` 2) (s + 1)
+      | otherwise = (x, s)
+
+-- Tonelli-Shanks algorithm: find x such that x^2 = n (mod p)
+-- Returns Just x if it exists, Nothing otherwise.
+tonelliShanks :: Integer -> Integer -> Maybe Integer
+tonelliShanks n p
+  | legendre n p /= 1 = Nothing   -- no square root
+  | p == 2            = Just n    -- special case
+  | otherwise         = Just x
+  where
+    (q, s) = factorOutTwos (p - 1)
+
+    -- find z which is a quadratic non-residue mod p
+    z = head [z' | z' <- [2..p-1], legendre z' p == p - 1]
+
+    m0 = s
+    c0 = modExp z q p
+    t0 = modExp n q p
+    r0 = modExp n ((q + 1) `div` 2) p
+
+    (x, _, _) = loop m0 c0 t0 r0
+
+    loop :: Integer -> Integer -> Integer -> Integer -> (Integer, Integer, Integer)
+    loop m c t r
+      | t == 0 = (0, c, t)
+      | t == 1 = (r, c, t)
+      | otherwise =
+          let
+            i = smallestI 0 t m
+            b = modExp c (2 ^ (m - i - 1)) p
+            m' = i
+            c' = modMul b b p
+            t' = modMul t c' p
+            r' = modMul r b p
+          in loop m' c' t' r'
+
+    -- find smallest i (0 <= i < m) such that t^(2^i) = 1 mod p
+    smallestI i t m
+      | i >= m = error "no valid i found"
+      | modExp t (2 ^ i) p == 1 = i
+      | otherwise = smallestI (i + 1) t m