misc/haskell/Miller-Rabin.hs

59 lines
1.5 KiB
Haskell
Raw Permalink Normal View History

2018-08-05 18:53:07 +02:00
#!/usr/bin/env nix-script
#!>haskell
#! haskell | mtl random
import Control.Monad (replicateM)
import Control.Monad.State (State, state, evalState)
import System.Random (Random, StdGen, newStdGen, randomR)
import Data.List (genericLength)
-- Modular integral exponentiation
(^%) :: Integral a => a -> a -> a -> a
(^%) m b k = ex b k 1 where
ex a k s
| k == 0 = s
| k `mod` 2 == 0 = ex (a * a `mod` m) (k `div` 2) s
| otherwise = ex (a * a `mod` m) (k `div` 2) (s * a `mod` m)
-- Deconstruct an integral n into two
-- numbers d, r such as d*2^r = n
decons :: Integral a => a -> (a, a)
decons n = (exp quots, rest quots)
where
rest = head . dropWhile even
exp = genericLength . takeWhile even
quots = iterate (`quot` 2) n
-- Test a number given a witness
test :: Integer -> Integer -> Bool
test n a
| n < 2 || even n = False
| b0 == 1 || b0 == pred n = True
| otherwise = iter (tail b)
where
(r, d) = decons (pred n)
b0 = (n ^% a) d
b = take (fromIntegral r) (iterate (n ^% 2) b0)
iter [] = False
iter (x:xs)
| x == 1 = False
| x == pred n = True
| otherwise = iter xs
-- Check whether n is probably a prime number
-- with un uncertainty of 2^-k
isPrime :: Integer -> Int -> State StdGen Bool
isPrime n k = all (test n) <$> replicateM k a
where a = state (randomR (2, pred n))
main :: IO ()
main = do
gen <- newStdGen
let test n = evalState (isPrime n 64) gen
print (test 9917340299)