Created
November 25, 2017 18:30
-
-
Save dpiponi/fc2817d75cbd41fbc0311c338cd2d591 to your computer and use it in GitHub Desktop.
Primality test based on Chebyshev polynomials
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import Data.Bits | |
| import Control.Monad | |
| type Z = Integer | |
| -- Find smallest power of two >= given integer. | |
| -- Sadly it's not convenient using the usual interface to Integer | |
| -- Got exceptions when using Data.Bits.Bitwise | |
| suitablePower :: Z -> Int | |
| suitablePower 0 = -1 | |
| suitablePower n = suitablePower' n 0 where | |
| suitablePower' n i = if shiftL 1 i >= n | |
| then i | |
| else suitablePower' n (i+1) | |
| extendedGCD :: Z -> Z -> (Z, Z, Z) | |
| extendedGCD a 0 = (a, 1, 0) | |
| extendedGCD a b = let (q, r) = a `quotRem` b | |
| (g, s, t) = extendedGCD b r | |
| in (abs g, t, s - q * t) | |
| zipWithDefault f [] [] = [] | |
| zipWithDefault f [] as = as | |
| zipWithDefault f as [] = as | |
| zipWithDefault f (a : as) (b : bs) = f a b : zipWithDefault f as bs | |
| convolveWith :: Ring a -> [a] -> [a] -> [a] | |
| convolveWith ring [a] bs = map (times ring a) bs | |
| convolveWith ring (a : as) bs = | |
| zipWithDefault (plus ring) (map (times ring a) bs) | |
| (zero ring : convolveWith ring as bs) | |
| -- Rings | |
| data Ring a = R { plus :: a -> a -> a, times :: a -> a -> a, zero :: a, one :: a, num :: Z -> a } | |
| -- From Monoid source | |
| power ring x0 y0 | y0 < 0 = error "Negative exponent" | |
| | y0 == 0 = one ring | |
| | otherwise = f x0 y0 | |
| where f x y | even y = f (times ring x x) (y `quot` 2) | |
| | y == 1 = x | |
| | otherwise = g (times ring x x) ((y - 1) `quot` 2) x | |
| g x y z | even y = g (times ring x x) (y `quot` 2) z | |
| | y == 1 = times ring x z | |
| | otherwise = g (times ring x x) ((y - 1) `quot` 2) (times ring x z) | |
| -- The Montgomery representation | |
| -- Based on https://en.wikipedia.org/wiki/Montgomery_modular_multiplication | |
| newtype Montgomery = N Z deriving Show | |
| toMontgomery :: Int -> Z -> Z -> Montgomery | |
| toMontgomery i n a = N $ (shiftL a i) `mod` n | |
| fromMontgomery :: Z -> Z -> Montgomery -> Z | |
| fromMontgomery invs n (N a) = (invs*a) `mod` n | |
| montgomeryReduce :: Int -> Z -> Z -> Z -> Z -> Z | |
| montgomeryReduce i mask n invn t = | |
| let tm = ((t .&. mask)*invn) .&. mask | |
| tt = shiftR (t + tm*n) i | |
| in if tt >= n then tt-n else tt | |
| montgomeryPlus :: Z -> Montgomery -> Montgomery -> Montgomery | |
| montgomeryPlus n (N a) (N b) = let u = a+b in N (if u >= n then u-n else u) | |
| montgomeryMultiply :: Int -> Z -> Z -> Z -> Montgomery -> Montgomery -> Montgomery | |
| montgomeryMultiply i mask n invn (N a) (N b) = N $ montgomeryReduce i mask n invn (a*b) | |
| -- Return pair consisting of | |
| -- 1. The ring structure for montogomery arithmetic | |
| -- 2. The function to map back to the ordinary integers | |
| montgomery :: Z -> (Ring Montgomery, Montgomery -> Z) | |
| montgomery n = | |
| let i = suitablePower n | |
| s = shiftL 1 i | |
| mask = s-1 | |
| (_, a, b) = extendedGCD n s | |
| in (R { plus = montgomeryPlus n, | |
| times = montgomeryMultiply i mask n ((-a) .&. mask), | |
| zero = toMontgomery i n 0, | |
| one = toMontgomery i n 1, | |
| num = toMontgomery i n}, | |
| fromMontgomery (b `mod` n) n) | |
| -- Polynomial rings | |
| newtype Polynomial a = P [a] deriving Show | |
| reducedPolyMultiply :: Ring a -> a -> | |
| Z -> Polynomial a -> Polynomial a -> Polynomial a | |
| reducedPolyMultiply ring zero r (P a) (P b) = P $ let ps = convolveWith ring a b | |
| in uncurry (zipWithDefault (plus ring)) (splitAt (fromIntegral r) ps) | |
| polyPlus :: (a -> a -> a) -> Polynomial a -> Polynomial a -> Polynomial a | |
| polyPlus (+) (P a) (P b) = P $ zipWithDefault (+) a b | |
| x :: Ring a -> Polynomial a | |
| x ring = P [num ring 0, num ring 1] | |
| makePolynomials :: Ring a -> Z -> Ring (Polynomial a) | |
| makePolynomials ring r = R { | |
| plus = polyPlus (plus ring), | |
| times = reducedPolyMultiply ring (num ring 0) r, | |
| zero = P [], | |
| one = P [num ring 1], | |
| num = \i -> P [num ring i] | |
| } | |
| -- Matrix rings | |
| data Matrix a = M a a a a deriving Show | |
| matrixMultiply (+) (*) (M a b c d) (M e f g h) = | |
| M ((a*e)+(b*g)) ((a*f)+(b*h)) | |
| ((c*e)+(d*g)) ((c*f)+(d*h)) | |
| get :: Ring a -> Ring (Polynomial a) -> Matrix (Polynomial a) -> Polynomial a | |
| get ring polynomials (M a b _ _) = plus polynomials (times polynomials a (x ring)) b | |
| makeMatrices :: Ring a -> Ring (Matrix a) | |
| makeMatrices ring = R { | |
| plus = undefined, | |
| times = matrixMultiply (plus ring) (times ring), | |
| zero = M (zero ring) (zero ring) | |
| (zero ring) (zero ring), | |
| one = M (one ring) (zero ring) | |
| (zero ring) (one ring), | |
| num = \i -> let ii = num ring i | |
| zero = num ring 0 | |
| in M ii zero | |
| zero ii | |
| } | |
| -- Compute Chebyshev polynomials using | |
| -- https://mathoverflow.net/questions/286626/is-there-an-explicit-expression-for-chebyshev-polynomials-modulo-xr-1/286639#286639 | |
| generator :: Ring a -> Matrix (Polynomial a) | |
| generator ring = M (P [num ring 0, num ring 2]) (P [num ring (-1)]) | |
| (P [num ring 1]) (P [num ring 0]) | |
| -- Compute mth Chebyshev polynomial modulo (n, x^r-1) | |
| chebyshev :: Ring a -> Ring (Polynomial a) -> Ring (Matrix (Polynomial a)) -> Z -> Polynomial a | |
| chebyshev base polynomials matrices m = | |
| let chebyshevMatrix = power matrices (generator base) (m-1) | |
| in get base polynomials chebyshevMatrix | |
| chebyshevTest :: Z -> Z -> Z -> [Z] | |
| chebyshevTest m r n = | |
| let (base, extract) = montgomery n | |
| polynomials = makePolynomials base r | |
| matrices = makeMatrices polynomials | |
| P as = chebyshev base polynomials matrices m | |
| in map extract as | |
| -- Borrowed from the web somewhere... :-) | |
| primes :: [Z] | |
| primes = sieve [2..] | |
| where sieve (p : xs) = p : sieve [x | x <- xs, x `mod` p /= 0] | |
| -- Check if list is all zeros except for a 1 at given position. | |
| allZeros [] = True | |
| allZeros (0 : as) = allZeros as | |
| allZeros _ = False | |
| allZerosBut 0 (1 : as) = allZeros as | |
| allZerosBut i (0 : as) = allZerosBut (i-1) as | |
| allZerosBut _ _ = False | |
| -- Primality test based on | |
| -- https://mathoverflow.net/questions/286304/chebyshev-polynomials-of-the-first-kind-and-primality-testing | |
| -- Way slower than I'd like. | |
| -- Don't know if this because | |
| -- 1. my implementation is poor | |
| -- 2. the methods that people actually use are faster in practice. | |
| -- But this implementation, if the conjecture is proved, would | |
| -- have a guaranteed deterministic running time polynomial | |
| -- in the size of the queried number. | |
| primeq :: Z -> Bool | |
| primeq n = prime' (tail primes) n where | |
| prime' (r:rs) n | mod n r == 0 = False | |
| | mod (n*n) r /= 1 = allZerosBut (n `mod` r) $ chebyshevTest n r n | |
| | otherwise = prime' rs n | |
| main = do | |
| print "Start" | |
| -- print $ primeq (2^86243-1) -- Made laptop reboot! | |
| -- Find Mersenne primes | |
| forM [2..] $ \i -> | |
| when (primeq (2^i-1)) $ print i |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Code doesn't work for primes 2,3 and 5.