-- |
-- Module      : Data.Numbers.Primes
-- Copyright   : Sebastian Fischer
-- License     : PublicDomain
-- 
-- Maintainer  : Sebastian Fischer (sebf@informatik.uni-kiel.de)
-- Stability   : experimental
-- Portability : portable
-- 
-- This Haskell library provides an efficient lazy wheel sieve for
-- prime generation ispired by "Lazy wheel sieves and spirals of
-- primes" [1] by Colin Runciman and "The Genuine Sieve of
-- Eratosthenes" [2] by Melissa O'Neil.
-- 
-- [1]: <http://www.cs.york.ac.uk/ftpdir/pub/colin/jfp97lw.ps.gz>
-- 
-- [2]: <http://www.cs.hmc.edu/~oneill/papers/Sieve-JFP.pdf>
-- 
module Data.Numbers.Primes ( primes, wheelSieve ) where

-- | 
-- This global constant is an infinite list of prime numbers. It is
-- generated by a lazy wheel sieve and shared among different
-- applications. If you are concerned about the memory requirements of
-- sharing many primes you can call the function @wheelSieve@
-- directly.
-- 
primes :: [Integer]
primes = wheelSieve 6

-- |
-- This function returns an infinite list of prime numbers by sieving
-- with a wheel that cancels the multiples of the first @n@ primes
-- where @n@ is the argument given to @wheelSieve@. Don't use too
-- large wheels because computing them is more expensive than
-- sieving. The number @6@ is a good value to pass to this function.
-- 
wheelSieve :: Int        -- ^ number of primes canceled by the wheel
           -> [Integer]  -- ^ infinite list of primes
wheelSieve k = reverse ps ++ sieve (spin p (cycle ns)) Empty 
 where (p:ps,ns)     = wheel k
       spin n (x:xs) = n : spin (n+x) xs


-- Auxiliary Definitions
------------------------------------------------------------------------------

-- Sieves a list of prime candidates using a lazy priority queue.
--
sieve :: [Integer] -> Queue -> [Integer] 
sieve (n:ns) Empty = n : sieve ns (enqueue (map (n*) (n:ns)) Empty) 
sieve (n:ns) queue 
  | m == n         = sieve ns (enqueue ms q) 
  | m < n          = sieve (n:ns) (enqueue ms q) 
  | otherwise      = n : sieve ns (enqueue (map (n*) (n:ns)) queue) 
 where (m:ms,q) = dequeue queue 

-- 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])

-- Computes a wheel that cancels the multiples of the given number
-- (plus 1) of primes.
--
-- For example:
--
-- wheel 0 = ([2],[1]) 
-- wheel 1 = ([3,2],[2]) 
-- 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 n = iterate next ([2],[1]) !! n

next :: Wheel -> Wheel
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]
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


-- We use a special version of priority queues implemented as "pairing
-- heaps" ((see "Purely functional datastructures by Chris Okasaki).
--
-- The queue stores non-empty lists of multiples; the first element is
-- used as priority.
--
data Queue = Empty | Fork [Integer] [Queue] 

enqueue :: [Integer] -> Queue -> Queue
enqueue ns = merge (Fork ns []) 

merge :: Queue -> Queue -> Queue
merge Empty y = y; merge x Empty = x 
merge x y | prio x <= prio y = join x y 
          | otherwise        = join y x 
 where prio (Fork (n:_) _)   = n 
       join (Fork ns qs) q   = Fork ns (q:qs) 

dequeue :: Queue -> ([Integer], Queue)
dequeue (Fork ns qs) = (ns,mergeAll qs) 

mergeAll :: [Queue] -> Queue
mergeAll [] = Empty; mergeAll [x] = x 
mergeAll (x:y:qs) = merge (merge x y) (mergeAll qs)