52 lines
1.5 KiB
Haskell
52 lines
1.5 KiB
Haskell
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
|