59 lines
1.5 KiB
Haskell
59 lines
1.5 KiB
Haskell
|
#!/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)
|