Skip to content

Instantly share code, notes, and snippets.

@dendisuhubdy
Forked from dpiponi/main.hs
Created November 25, 2017 23:03
Show Gist options
  • Save dendisuhubdy/d4a342c4be4d4e5f5af30666d9df1099 to your computer and use it in GitHub Desktop.
Save dendisuhubdy/d4a342c4be4d4e5f5af30666d9df1099 to your computer and use it in GitHub Desktop.
Primality test based on Chebyshev polynomials
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