Primality test based on Chebyshev polynomials

main.hs

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