Noob experiments on the prime generator - see the paper from Melissa E. O’Neill

prime-tests.hs

#!/usr/bin/env stack
-- stack --resolver lts-11.14 --install-ghc runghc --package data-ordlist --package PSQueue-1.1 --package pqueue-1.4.1.1

import Text.Printf
import Data.List.Ordered ( minus )
import qualified Data.Map as Map
import qualified Data.PSQueue as PSQ
import qualified Data.PQueue.Min as PQ
import System.CPUTime

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

primes2 :: [ Int ]
primes2 = 2 : 3 : nextPrime primes2 [ 4 .. ]
    where
    nextPrime restPrimes integrals =
        let ( p : ps ) = restPrimes
            notMultples = filter ( \x -> x `mod` p /= 0 ) integrals
        in head notMultples : nextPrime ps ( tail notMultples )

primes3 :: [ Int ]
primes3 = sieve [ 2 .. ]
    where
    sieve [] = []
    sieve ( p:xs ) = p : sieve( xs `minus` [ p*p, p*p+p .. ] )

trialDivisionPrimes = 2 : [ x | x <- [ 3 .. ], isPrime x ]
    where
    isPrime x = all ( \p -> x `mod` p > 0 ) ( candidateFactors x )
    candidateFactors x = takeWhile ( \p -> p*p <= x ) trialDivisionPrimes

primes = 2:([3..] `minus` composites)
    where
    composites = union [multiples p | p <- primes]
    multiples n = map (n*) [n..]
    (x:xs) `minus` (y:ys)
        | x < y = x:(xs `minus` (y:ys))
        | x == y = xs `minus` ys
        | x > y = (x:xs) `minus` ys

union = foldr merge []
    where
    merge (x:xs) ys = x:merge' xs ys
    merge' (x:xs) (y:ys)
        | x < y = x:merge' xs (y:ys)
        | x == y = x:merge' xs ys
        | x > y = y:merge' (x:xs) ys

candidates = [3,5..]

mapPrime = sieve [ 2 .. ]
    where
    sieve xs = sieve' xs Map.empty
        where
        sieve' [] table = []
        sieve' (x:xs) table =
            case Map.lookup x table of
                Nothing -> x : sieve' xs ( Map.insert ( x*x ) [ x ] table )
                Just facts -> sieve' xs ( foldl reinsert ( Map.delete x table ) facts )
                    where
                    reinsert table prime = Map.insertWith (++) ( x + prime ) [ prime ] table

psqPrime = 2 : sieve candidates
    where
    sieve [] = []
    sieve (x:xs) = x : sieve' xs ( insertPrime x xs PSQ.empty )
        where
        insertPrime p xs table = PSQ.insert ( map ( * p ) xs ) ( p*p ) table
        sieve' [] table = []
        sieve' (x:xs) table
            | nextComposite <= x = sieve' xs ( adjust table )
            | otherwise = x : sieve' xs ( insertPrime x xs table )
            where
            nextComposite = case PSQ.findMin table of
                Just minBinding -> PSQ.prio minBinding
                Nothing -> error "Empty PSQ!"
            adjust table
                | n <= x = adjust ( PSQ.insert ns n' . PSQ.deleteMin $ table )
                | otherwise = table
                where
                ( n':ns, n ) = case PSQ.findMin table of
                    Just ( k PSQ.:-> v ) -> ( k, v )
                    Nothing -> error "Empty PSQ!"

-- adding the simple wheel will make this quite fast!
-- pqPrime = 2 : 3 : 5 : 7 : sieve ( spin wheel 11 )
pqPrime = 2 : sieve candidates
    where
    sieve [] = []
    sieve (x:xs) = x : sieve' xs ( insertPrime x xs PQ.empty )
        where
        insertPrime p xs table = PQ.insert ( p*p, map ( * p ) xs )  table
        sieve' [] table = []
        sieve' (x:xs) table
            | nextComposite <= x = sieve' xs ( adjust table )
            | otherwise = x : sieve' xs ( insertPrime x xs table )
            where
            ( nextComposite, _ ) = PQ.findMin table
            adjust table
                | n <= x = adjust ( PQ.insert ( n', ns ) . PQ.deleteMin $ table )
                | otherwise = table
                where
                ( n, n':ns ) = PQ.findMin table

data PriorityQ k v = Lf
                   | Br {-# UNPACK #-} !k v !(PriorityQ k v) !(PriorityQ k v)
               deriving (Eq, Ord, Read, Show)

emptyPQ :: PriorityQ k v
emptyPQ = Lf

isEmptyPQ :: PriorityQ k v -> Bool
isEmptyPQ Lf  = True
isEmptyPQ _   = False

minKeyValuePQ :: PriorityQ k v -> (k, v)
minKeyValuePQ (Br k v _ _)    = (k,v)
minKeyValuePQ _               = error "Empty heap!"

minKeyPQ :: PriorityQ k v -> k
minKeyPQ (Br k v _ _)         = k
minKeyPQ _                    = error "Empty heap!"

minValuePQ :: PriorityQ k v -> v
minValuePQ (Br k v _ _)       = v
minValuePQ _                  = error "Empty heap!"

insertPQ :: Ord k => k -> v -> PriorityQ k v -> PriorityQ k v
insertPQ wk wv (Br vk vv t1 t2)
               | wk <= vk   = Br wk wv (insertPQ vk vv t2) t1
               | otherwise  = Br vk vv (insertPQ wk wv t2) t1
insertPQ wk wv Lf             = Br wk wv Lf Lf

siftdown :: Ord k => k -> v -> PriorityQ k v -> PriorityQ k v -> PriorityQ k v
siftdown wk wv Lf _             = Br wk wv Lf Lf
siftdown wk wv (t @ (Br vk vv _ _)) Lf
    | wk <= vk                  = Br wk wv t Lf
    | otherwise                 = Br vk vv (Br wk wv Lf Lf) Lf
siftdown wk wv (t1 @ (Br vk1 vv1 p1 q1)) (t2 @ (Br vk2 vv2 p2 q2))
    | wk <= vk1 && wk <= vk2    = Br wk wv t1 t2
    | vk1 <= vk2                = Br vk1 vv1 (siftdown wk wv p1 q1) t2
    | otherwise                 = Br vk2 vv2 t1 (siftdown wk wv p2 q2)

deleteMinAndInsertPQ :: Ord k => k -> v -> PriorityQ k v -> PriorityQ k v
deleteMinAndInsertPQ wk wv Lf             = error "Empty PriorityQ"
deleteMinAndInsertPQ wk wv (Br _ _ t1 t2) = siftdown wk wv t1 t2

leftrem :: PriorityQ k v -> (k, v, PriorityQ k v)
leftrem (Br vk vv Lf Lf) = (vk, vv, Lf)
leftrem (Br vk vv t1 t2) = (wk, wv, Br vk vv t2 t) where
    (wk, wv, t) = leftrem t1
leftrem _                = error "Empty heap!"

deleteMinPQ :: Ord k => PriorityQ k v -> PriorityQ k v
deleteMinPQ (Br vk vv Lf _) = Lf
deleteMinPQ (Br vk vv t1 t2) = siftdown wk wv t2 t where
    (wk,wv,t) = leftrem t1
deleteMinPQ _ = error "Empty heap!"

-- A hybrid of Priority Queues and regular queues.  It allows a priority
-- queue to have a feeder queue, filled with items that come in an
-- increasing order.  By keeping the feed for the queue separate, we
-- avoid needlessly filling an O(log n) data structure with data that
-- it won't need for a long time.

type HybridQ k v = (PriorityQ k v, [(k,v)])

initHQ :: PriorityQ k v -> [(k,v)] -> HybridQ k v
initHQ pq feeder = (pq, feeder)

insertHQ :: (Ord k) => k -> v -> HybridQ k v -> HybridQ k v
insertHQ k v (pq, q) = (insertPQ k v pq, q)

deleteMinAndInsertHQ :: (Ord k) => k -> v -> HybridQ k v -> HybridQ k v
deleteMinAndInsertHQ k v (pq, q) = postRemoveHQ(deleteMinAndInsertPQ k v pq, q)
    where
        postRemoveHQ mq@(pq, []) = mq
        postRemoveHQ mq@(pq, (qk,qv) : qs)
            | qk < minKeyPQ pq = (insertPQ qk qv pq, qs)
            | otherwise        = mq

minKeyHQ      :: HybridQ k v -> k
minKeyHQ (pq, q) = minKeyPQ pq

minKeyValueHQ :: HybridQ k v -> (k, v)
minKeyValueHQ (pq, q) = minKeyValuePQ pq

wheel :: Integral a => [a]
wheel = 2:4:2:4:6:2:6:4:2:4:6:6:2:6:4:2:6:4:6:8:4:2:4:2:4:8:6:4:6:2:4:6:2:6:6:4:2:4:6:2:6:4:2:4:2:10:2:10:wheel

spin (x:xs) n = n : spin xs (n+x)

originalONeilPrimes = 2 : 3 : 5 : 7 : sieve ( spin wheel 11 )
    where
    sieve [] = []
    sieve (x:xs) = x : sieve' xs (insertprime x xs emptyPQ)
        where
        insertprime p xs table = insertPQ (p*p) (map (* p) xs) table
        sieve' [] table = []
        sieve' (x:xs) table
            | nextComposite <= x = sieve' xs (adjust table)
            | otherwise = x : sieve' xs (insertprime x xs table)
            where
            nextComposite = minKeyPQ table
            adjust table
                | n <= x = adjust (deleteMinAndInsertPQ n' ns table)
                | otherwise = table
                where
                (n, n':ns) = minKeyValuePQ table

birdPrimes = 2 : ( [ 3 .. ] `minus` composites )
    where
    composites = union [ multiples p | p <- birdPrimes ]
    multiples n = map ( n* ) [ n .. ]
    ( x:xs ) `minus` ( y:ys )
        | x < y = x : ( xs `minus` ( y:ys ) )
        | x == y = xs `minus` ys
        | x > y =  ( x:xs ) `minus` ys
    union = foldr merge []
        where
        merge ( x:xs ) ys = x:merge' xs ys
        merge' ( x:xs ) ( y:ys )
            | x < y = x:merge' xs ( y:ys )
            | x == y = x:merge' xs ys
            | x > y = y:merge' ( x:xs ) ys

simpleTime label v = do
    begin <- getCPUTime

    print v

    end <- getCPUTime

    let diff = ( fromIntegral ( end - begin ) ) / ( 10^12 )
    printf "%s costs: %0.3f sec\n" label ( diff :: Double )

takeTerminate = last . take 50000

main = do
    simpleTime "originalONeilPrimes" ( takeTerminate originalONeilPrimes )
    simpleTime "mapPrime" ( takeTerminate mapPrime )
    simpleTime "pqPrime" ( takeTerminate pqPrime )
    simpleTime "psqPrime" ( takeTerminate psqPrime )
    simpleTime "trialDivisionPrimes" ( takeTerminate trialDivisionPrimes )
    simpleTime "birdPrimes" ( takeTerminate birdPrimes )