diff --git a/Data/Numbers/Primes.hs b/Data/Numbers/Primes.hs --- a/Data/Numbers/Primes.hs +++ b/Data/Numbers/Primes.hs @@ -14,17 +14,26 @@ -- /The Genuine Sieve of Eratosthenes/ by Melissa O'Neil -- (). -- -module Data.Numbers.Primes ( primes, wheelSieve ) where +module Data.Numbers.Primes ( + primes, wheelSieve, + + isPrime, primeFactors + + ) where + -- | -- This global constant is an infinite list of prime numbers. It is -- generated by a lazy wheel sieve and shared across the whole program -- run. If you are concerned about the memory requirements of sharing -- many primes you can call the function @wheelSieve@ directly. -- -primes :: [Integer] +primes :: Integral int => [int] primes = wheelSieve 6 +{-# SPECIALISE primes :: [Int] #-} +{-# SPECIALISE primes :: [Integer] #-} + -- | -- This function returns an infinite list of prime numbers by sieving -- with a wheel that cancels the multiples of the first @n@ primes @@ -33,11 +42,47 @@ -- function. Larger wheels improve the run time at the cost of higher -- memory requirements. -- -wheelSieve :: Int -- ^ number of primes canceled by the wheel - -> [Integer] -- ^ infinite list of primes +wheelSieve :: Integral int + => Int -- ^ number of primes canceled by the wheel + -> [int] -- ^ infinite list of primes wheelSieve k = reverse ps ++ map head (sieve p (cycle ns)) where (p:ps,ns) = wheel k +{-# SPECIALISE wheelSieve :: Int -> [Int] #-} +{-# SPECIALISE wheelSieve :: Int -> [Integer] #-} + +-- | +-- Checks whether a given positive number is prime. +-- +-- This function uses trial division to check for divisibility with +-- all primes below the square root of the given number. It is +-- impractical for numbers with a very large smallest prime factor. +-- +isPrime :: Integral int => int -> Bool +isPrime n = primeFactors n == [n] + +{-# SPECIALISE isPrime :: Int -> Bool #-} +{-# SPECIALISE isPrime :: Integer -> Bool #-} + +-- | +-- Yields the sorted list of prime factors of the given positive +-- number. +-- +-- This function uses trial division and is impractical for numbers +-- with very large prime factors. +-- +primeFactors :: Integral int => int -> [int] +primeFactors n = factors n (wheelSieve 6) + where + factors 1 _ = [] + factors m (p:ps) | m < p*p = [m] + | r == 0 = p : factors q (p:ps) + | otherwise = factors m ps + where (q,r) = quotRem m p + +{-# SPECIALISE primeFactors :: Int -> [Int] #-} +{-# SPECIALISE primeFactors :: Integer -> [Integer] #-} + -- Auxiliary Definitions ------------------------------------------------------------------------------ @@ -53,60 +98,81 @@ -- that need to be cancelled, one can multiply all elements of the -- list with its head. -- -sieve :: Integer -> [Integer] -> [[Integer]] +sieve :: (Ord int, Num int) => int -> [int] -> [[int]] sieve p ns@(m:ms) = spin p ns : sieveComps (p+m) ms (composites p ns) +{-# SPECIALISE sieve :: Int -> [Int] -> [[Int]] #-} +{-# SPECIALISE sieve :: Integer -> [Integer] -> [[Integer]] #-} + -- Composites are stored in increasing order in a priority queue. The -- queue has an associated feeder which is used to avoid filling it -- with entries that will only be used again much later. -- -type Composites = (Queue,[[Integer]]) +type Composites int = (Queue int,[[int]]) -- The feeder is computed from the result of a call to 'sieve'. -- -composites :: Integer -> [Integer] -> Composites +composites :: (Ord int, Num int) => int -> [int] -> Composites int composites p ns = (Empty, map comps (spin p ns : sieve p ns)) where comps xs@(x:_) = map (x*) xs +{-# SPECIALISE composites :: Int -> [Int] -> Composites Int #-} +{-# SPECIALISE composites :: Integer -> [Integer] -> Composites Integer #-} + -- We can split all composites into the next and remaining -- composites. We use the feeder when appropriate and discard equal -- entries to not return a composite twice. -- -splitComposites :: Composites -> (Integer,Composites) +splitComposites :: Ord int => Composites int -> (int,Composites int) splitComposites (Empty, xs:xss) = splitComposites (Fork xs [], xss) splitComposites (queue, xss@((x:xs):yss)) | x < z = (x, discard x (enqueue xs queue, yss)) | otherwise = (z, discard z (enqueue zs queue', xss)) where (z:zs,queue') = dequeue queue +{-# SPECIALISE splitComposites :: Composites Int -> (Int,Composites Int) #-} +{-# SPECIALISE + splitComposites :: Composites Integer -> (Integer,Composites Integer) #-} + -- Drops all occurrences of the given element. -- -discard :: Integer -> Composites -> Composites +discard :: Ord int => int -> Composites int -> Composites int discard n ns | n == m = discard n ms | otherwise = ns where (m,ms) = splitComposites ns +{-# SPECIALISE discard :: Int -> Composites Int -> Composites Int #-} +{-# SPECIALISE + discard :: Integer -> Composites Integer -> Composites Integer #-} + -- This is the actual sieve. It discards candidates that are -- composites and yields lists which start with a prime and contain -- all factors of the composites that need to be dropped. -- -sieveComps :: Integer -> [Integer] -> Composites -> [[Integer]] +sieveComps :: (Ord int, Num int) => int -> [int] -> Composites int -> [[int]] sieveComps cand ns@(m:ms) xs | cand == comp = sieveComps (cand+m) ms ys | cand < comp = spin cand ns : sieveComps (cand+m) ms xs | otherwise = sieveComps cand ns ys where (comp,ys) = splitComposites xs +{-# SPECIALISE sieveComps :: Int -> [Int] -> Composites Int -> [[Int]] #-} +{-# SPECIALISE + sieveComps :: Integer -> [Integer] -> Composites Integer -> [[Integer]] #-} + -- This function computes factors of composites of primes by spinning -- a wheel. -- -spin :: Integer -> [Integer] -> [Integer] +spin :: Num int => int -> [int] -> [int] spin x (y:ys) = x : spin (x+y) ys +{-# SPECIALISE spin :: Int -> [Int] -> [Int] #-} +{-# SPECIALISE spin :: Integer -> [Integer] -> [Integer] #-} + -- A wheel consists of a list of primes whose multiples are canceled -- and the actual wheel that is rolled for canceling. -- -type Wheel = ([Integer],[Integer]) +type Wheel int = ([int],[int]) -- Computes a wheel that cancels the multiples of the given number -- (plus 1) of primes. @@ -118,33 +184,46 @@ -- wheel 2 = ([5,3,2],[2,4]) -- wheel 3 = ([7,5,3,2],[4,2,4,2,4,6,2,6]) -- -wheel :: Int -> Wheel +wheel :: Integral int => Int -> Wheel int wheel n = iterate next ([2],[1]) !! n -next :: Wheel -> Wheel +{-# SPECIALISE wheel :: Int -> Wheel Int #-} +{-# SPECIALISE wheel :: Int -> Wheel Integer #-} + +next :: Integral int => Wheel int -> Wheel int next (ps@(p:_),xs) = (py:ps,cancel (product ps) p py ys) where (y:ys) = cycle xs py = p + y -cancel :: Integer -> Integer -> Integer -> [Integer] -> [Integer] +{-# SPECIALISE next :: Wheel Int -> Wheel Int #-} +{-# SPECIALISE next :: Wheel Integer -> Wheel Integer #-} + +cancel :: Integral int => int -> int -> int -> [int] -> [int] cancel 0 _ _ _ = [] cancel m p n (x:ys@(y:zs)) | nx `mod` p > 0 = x : cancel (m-x) p nx ys | otherwise = cancel m p n (x+y:zs) where nx = n + x +{-# SPECIALISE cancel :: Int -> Int -> Int -> [Int] -> [Int] #-} +{-# SPECIALISE + cancel :: Integer -> Integer -> Integer -> [Integer] -> [Integer] #-} + -- We use a special version of priority queues implemented as /pairing/ -- /heaps/ (see /Purely Functional Data Structures/ by Chris Okasaki). -- -- The queue stores non-empty lists of composites; the first element -- is used as priority. -- -data Queue = Empty | Fork [Integer] [Queue] +data Queue int = Empty | Fork [int] [Queue int] -enqueue :: [Integer] -> Queue -> Queue +enqueue :: Ord int => [int] -> Queue int -> Queue int enqueue ns = merge (Fork ns []) -merge :: Queue -> Queue -> Queue +{-# SPECIALISE enqueue :: [Int] -> Queue Int -> Queue Int #-} +{-# SPECIALISE enqueue :: [Integer] -> Queue Integer -> Queue Integer #-} + +merge :: Ord int => Queue int -> Queue int -> Queue int merge Empty y = y merge x Empty = x merge x y | prio x <= prio y = join x y @@ -152,10 +231,19 @@ where prio (Fork (n:_) _) = n join (Fork ns qs) q = Fork ns (q:qs) -dequeue :: Queue -> ([Integer], Queue) +{-# SPECIALISE merge :: Queue Int -> Queue Int -> Queue Int #-} +{-# SPECIALISE merge :: Queue Integer -> Queue Integer -> Queue Integer #-} + +dequeue :: Ord int => Queue int -> ([int], Queue int) dequeue (Fork ns qs) = (ns,mergeAll qs) -mergeAll :: [Queue] -> Queue +{-# SPECIALISE dequeue :: Queue Int -> ([Int], Queue Int) #-} +{-# SPECIALISE dequeue :: Queue Integer -> ([Integer], Queue Integer) #-} + +mergeAll :: Ord int => [Queue int] -> Queue int mergeAll [] = Empty mergeAll [x] = x mergeAll (x:y:qs) = merge (merge x y) (mergeAll qs) + +{-# SPECIALISE mergeAll :: [Queue Int] -> Queue Int #-} +{-# SPECIALISE mergeAll :: [Queue Integer] -> Queue Integer #-} diff --git a/memory.hs b/memory.hs new file mode 100644 --- /dev/null +++ b/memory.hs @@ -0,0 +1,10 @@ +-- ghc -O2 --make memory +-- ./memory 6 1000000 +RTS -s + +import System.Environment +import Data.Numbers.Primes + +main :: IO () +main = do [m,n] <- (map read . take 2) `fmap` getArgs + print $ (wheelSieve m :: [Int]) !! n + diff --git a/primes.cabal b/primes.cabal --- a/primes.cabal +++ b/primes.cabal @@ -1,5 +1,5 @@ Name: primes -Version: 0.1.1.1 +Version: 0.2.0.0 Cabal-Version: >= 1.6 Synopsis: Efficient, purely functional generation of prime numbers Description: @@ -19,7 +19,7 @@ Build-Type: Simple Stability: experimental -Extra-Source-Files: README +Extra-Source-Files: README, memory.hs, runtime.hs Library Build-Depends: base == 4.* diff --git a/runtime.hs b/runtime.hs new file mode 100644 --- /dev/null +++ b/runtime.hs @@ -0,0 +1,11 @@ +-- ghc -O2 --make runtime +-- ./runtime -u runtime.csv +-- barchart criterion runtime.csv + +import Criterion.Main +import Data.Numbers.Primes + +main :: IO () +main = defaultMain [ test size | size <- [5..15] ] + where test size = bench ("wheel"++show size) $ whnf comp size + comp size = (wheelSieve size :: [Int]) !! 10000