Prime numbers

Simple Prime Sieve

The following is an elegant (and highly inefficient) way to generate a list of all the prime numbers in the universe:

  primes :: [Integer]
  primes = sieve [2..]
    where sieve (p:xs) = p : sieve [x | x<-xs, x `mod` p /= 0]

Given an infinite list of prime numbers, we can implement primality tests and integer factorization:

  isPrime n = n > 1 && n == head (factorize n)

  primeFactors 1 = []
  primeFactors n = go n primes
     where
     go n ps@(p:pt)
        | p*p > n        = [n]
        | x `rem` p == 0 = p : go (n `quot` p) ps
        | otherwise      = go n pt

Simple Prime Sieve II

  primes :: [Integer]
  primes = 2:filter isPrime [3,5..]
     where
     isPrime n   = all (not . divides n) $ takeWhile (\p -> p*p <= n) primes
     divides n p = n `mod` p == 0

Prime Wheels

Notice in the above Prime Sieve II, only odd numbers are tested, because we know that all the even numbers (greater than 2) are composite. In effect, odd numbers, and not even numbers, are candidates for primality testing.

A prime wheel is a scheme to generate candidate numbers that are "pre-screened" so that they don't have certain predetermined divisors. For example, suppose we want candidates that are neither even nor divisible by 3. In that case, we need numbers of the form 6n + {1,5}.

primes :: [Integer]
primes = 2:3:5:filter isPrime wheel
 where
   -- these numbers are automatically not divisible by 2 or 3
   wheel = 7:11:map (6+) wheel
   -- don't bother to check for divisibility by 2 or 3
   ps = drop 2 primes
   isPrime n   = all (not . divides n) $ takeWhile (\p -> p*p <= n) ps
   divides n p = n `mod` p == 0

This generator runs slightly faster than Prime Sieve II above because it doesn't bother to perform prime testing on multiples of 2 or 3.

Here is why the scheme is called a prime wheel. Imagine that you had a wheel of circumference 6, and you are rolling that wheel along the number line. The wheel is marked along the edges to automatically tell you which numbers are candidates and which numbers to exclude. Specifically, multiples of 2, 3 or 6 are excluded, while numbers of the form 6n+1 and 6n+5 are candidates.

We can go further and exclude multiples of 5. To exclude multiples of 2, 3, and 5, our wheel has to increase in multiples of 30.

primes :: [Integer]
primes = 2:3:5:7:11:13:17:19:23:29:filter isPrime wheel
 where
   -- these numbers are automatically not divisible by 2, 3, or 5
   wheel = 31:37:41:43:47:49:53:59:map (30+) wheel
   -- don't bother to check for divisibility by 2, 3, or 5
   ps = drop 3 primes
   isPrime n   = all (not . divides n) $ takeWhile (\p -> p*p <= n) ps
   divides n p = n `mod` p == 0

This generator runs slightly faster than the (2,3) prime wheel because it doesn't bother to check multiples of 2, 3, or 5.

We can go even further and exclude multiples of 7, but this requires a much bigger wheel, and it provides only a very small additional speed-up. This wheel has a length of 210, and at this point we are probably well beyond the point of diminishing returns.

primes :: [Integer]
primes = initPrimes ++ filter isPrime wheel
 where
   initPrimes = [2,3,5,7,11,13,17,19,23,29,31,37,41,43,47,
                 53,59,61,67,71,73,79,83,89,97,
                 101,103,107,109,113,127,131,137,139,149,
                 151,157,163,167,173,179,181,191,193,197,199]
   -- the following numbers are automatically not divisible by 2, 3, 5, or 7
   wheel = [211,221,223,227,229,233,239,241,247,251,253,257,263,269,
            271,277,281,283,289,293,299,307,311,313,317,319,323,331,
            337,341,347,349,353,359,361,367,373,377,379,383,389,391,
            397,401,403,407,409,419] ++ map (210+) wheel
   -- don't bother to check for divisibility by 2, 3, 5, or 7
   ps = drop 4 primes
   isPrime n   = all (not . divides n) $ takeWhile (\p -> p*p <= n) ps
   divides n p = n `mod` p == 0

Implicit Heap

The following is a more efficient prime generator, implementing the sieve of Eratosthenes. See also the message thread around Re: Code and Perf. Data for Prime Finders.

data People a = VIP a (People a) | Crowd [a]

mergeP :: Ord a => People a -> People a -> People a
mergeP (VIP x xt) ys                    = VIP x $ mergeP xt ys
mergeP (Crowd xs) (Crowd ys)            = Crowd $ merge  xs ys
mergeP xs@(Crowd ~(x:xt)) ys@(VIP y yt) = case compare x y of
    LT -> VIP x $ mergeP (Crowd xt) ys
    EQ -> VIP x $ mergeP (Crowd xt) yt
    GT -> VIP y $ mergeP xs yt


merge :: Ord a => [a] -> [a] -> [a]
merge xs@(x:xt) ys@(y:yt) = case compare x y of
    LT -> x : merge xt ys
    EQ -> x : merge xt yt
    GT -> y : merge xs yt

diff xs@(x:xt) ys@(y:yt) = case compare x y of
    LT -> x : diff xt ys
    EQ ->     diff xt yt
    GT ->     diff xs yt

foldTree :: (a -> a -> a) -> [a] -> a
foldTree f ~(x:xs) = f x . foldTree f . pairs $ xs
    where pairs ~(x: ~(y:ys)) = f x y : pairs ys

primes, nonprimes :: [Integer]
primes    = 2:3:diff [5,7..] nonprimes
nonprimes = serve . foldTree mergeP . map multiples $ tail primes
    where
    multiples p = vip [p*k | k <- [p,p+2..]]

    vip (x:xs)       = VIP x $ Crowd xs
    serve (VIP x xs) = x:serve xs
    serve (Crowd xs) = xs

nonprimes effectively implements a heap, exploiting lazy evaluation.

Bitwise prime sieve

Count the number of prime below a given 'n'. Shows fast bitwise arrays, and an example of Template Haskell to defeat your enemies.

{-# OPTIONS -O2 -optc-O -XBangPatterns #-}
module Primes (nthPrime) where

import Control.Monad.ST
import Data.Array.ST
import Data.Array.Base
import System
import Control.Monad
import Data.Bits

nthPrime :: Int -> Int
nthPrime n = runST (sieve n)

sieve n = do
    a <- newArray (3,n) True :: ST s (STUArray s Int Bool)
    let cutoff = truncate (sqrt $ fromIntegral n) + 1
    go a n cutoff 3 1

go !a !m cutoff !n !c
    | n >= m    = return c
    | otherwise = do
        e <- unsafeRead a n
        if e then
            if n < cutoff then
                let loop !j
                        | j < m     = do
                            x <- unsafeRead a j
                            when x $ unsafeWrite a j False
                            loop (j+n)
                        | otherwise = go a m cutoff (n+2) (c+1)
                in loop ( if n < 46340 then n * n else n `shiftL` 1)
             else go a m cutoff (n+2) (c+1)
         else go a m cutoff (n+2) c

And places in a module:

{-# OPTIONS -fth #-}
import Primes

main = print $( let x = nthPrime 10000000 in [| x |] )

Run as:

$ ghc --make -o primes Main.hs
$ time ./primes
664579
./primes  0.00s user 0.01s system 228% cpu 0.003 total

Miller-Rabin Primality Test

find2km :: Integral a => a -> (a,a)
find2km n = f 0 n
    where 
        f k m
            | r == 1 = (k,m)
            | otherwise = f (k+1) q
            where (q,r) = quotRem m 2        
 
millerRabinPrimality :: Integer -> Integer -> Bool
millerRabinPrimality n a
    | a <= 1 || a >= n-1 = 
        error $ "millerRabinPrimality: a out of range (" 
              ++ show a ++ " for "++ show n ++ ")" 
    | n < 2 = False
    | even n = False
    | b0 == 1 || b0 == n' = True
    | otherwise = iter (tail b)
    where
        n' = n-1
        (k,m) = find2km n'
        b0 = powMod n a m
        b = take (fromIntegral k) $ iterate (squareMod n) b0
        iter [] = False
        iter (x:xs)
            | x == 1 = False
            | x == n' = True
            | otherwise = iter xs
 
pow' :: (Num a, Integral b) => (a -> a -> a) -> (a -> a) -> a -> b -> a
pow' _ _ _ 0 = 1
pow' mul sq x' n' = f x' n' 1
    where 
        f x n y
            | n == 1 = x `mul` y
            | r == 0 = f x2 q y
            | otherwise = f x2 q (x `mul` y)
            where
                (q,r) = quotRem n 2
                x2 = sq x
 
mulMod :: Integral a => a -> a -> a -> a
mulMod a b c = (b * c) `mod` a
squareMod :: Integral a => a -> a -> a
squareMod a b = (b * b) `rem` a
powMod :: Integral a => a -> a -> a -> a
powMod m = pow' (mulMod m) (squareMod m)