implement Tonelli Shanks

ModularArithmeticUtils.hs

@@ -0,0 +1,19 @@
+module ModularArithmeticUtils (modExp, modMul, legendre) where
+
+import Data.Bits (testBit, shiftR)
+
+-- Modular exponentiation: compute a^b mod m
+modExp :: Integer -> Integer -> Integer -> Integer
+modExp b 0 m = 1
+modExp b e m = t * modExp ((b * b) `mod` m) (shiftR e 1) m `mod` m
+            where t = if testBit e 0 then b `mod` m else 1
+
+-- Compute (a * b) mod m
+modMul :: Integer -> Integer -> Integer -> Integer
+modMul a b m = (a * b) `mod` m
+
+-- Compute the Legendre symbol (a/p)
+-- (a/p) = a^((p-1)/2) mod p
+-- (a/p) in {-1, 0, 1}
+legendre :: Integer -> Integer -> Integer
+legendre a p = modExp a ((p - 1) `div` 2) p

PollardPminus1.hs

@@ -1,7 +1,7 @@
 import Text.Read (readMaybe)
 import System.Exit (exitSuccess)
-import Data.Bits
 import Utils (askNumber)
+import ModularArithmeticUtils (modExp)
 
 main :: IO ()
 main = do
@@ -41,10 +41,3 @@ pollardP1WithParams n a b = go a 2
                else if d == n
                then Nothing -- Found n itself, try with different parameters
                else go newA (k + 1)
-
--- Modular exponentiation: compute a^b mod m
-modExp :: Integer -> Integer -> Integer -> Integer
-modExp b 0 m = 1
-modExp b e m = t * modExp ((b * b) `mod` m) (shiftR e 1) m `mod` m
-            where t = if testBit e 0 then b `mod` m else 1
-

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