Miller Rabin primality test

ModularArithmeticUtils.hs

@@ -1,4 +1,4 @@
-module ModularArithmeticUtils (modExp, modMul, legendre) where
+module ModularArithmeticUtils (modExp, modMul, legendre, factorOutTwos) where
 
 import Data.Bits (testBit, shiftR)
 
@@ -17,3 +17,11 @@ modMul a b m = (a * b) `mod` m
 -- (a/p) in {-1, 0, 1}
 legendre :: Integer -> Integer -> Integer
 legendre a p = modExp a ((p - 1) `div` 2) p
+
+-- 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)

README.md

@@ -10,6 +10,9 @@ file: `TonelliShanks.hs`
 ### Fermat's prime test
 file: `fermat-prime-test.hs`
 
+### Miller Rabin's prime test
+file: `miller-rabin-prime-test.hs`
+
 ## factorization algorithms
 
 ### Fermat's factorization

TonelliShanks.hs

@@ -1,14 +1,6 @@
 module TonelliShanks (tonelliShanks) where
 
-import ModularArithmeticUtils (modExp, modMul, legendre)
+import ModularArithmeticUtils (modExp, modMul, legendre, factorOutTwos)
-
--- 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.

miller-rabin-prime-test.hs

@@ -0,0 +1,28 @@
+import System.Random (randomRIO)
+import ModularArithmeticUtils (modExp, factorOutTwos)
+
+millerRabinWitness :: Integer -> Integer -> Integer -> Integer -> Bool
+millerRabinWitness n a d s =
+  let x0 = modExp a d n
+  in  x0 == 1 || x0 == n - 1 || loop x0 (s - 1)
+  where
+    loop _ 0 = False
+    loop x r =
+      let x' = (x * x) `mod` n
+      in  x' == n - 1 || loop x' (r - 1)
+
+millerRabin :: Integer -> Integer -> IO Bool
+millerRabin n k
+  | n < 2     = return False
+  | n == 2    = return True
+  | even n    = return False
+  | otherwise = do
+      let (d, s) = factorOutTwos (n - 1)
+      go k d s
+  where
+    go 0 _ _ = return True
+    go i d s = do
+      a <- randomRIO (2, n - 2)
+      if millerRabinWitness n a d s
+        then go (i - 1) d s
+        else return False