-- |
-- Module      : Data.Numbers.Primes
-- Copyright   : Sebastian Fischer
-- License     : BSD3
-- 
-- Maintainer  : Sebastian Fischer (sebf@informatik.uni-kiel.de)
-- Stability   : experimental
-- Portability : portable
-- 
-- This Haskell library provides an efficient lazy wheel sieve for
-- prime generation inspired by /Lazy wheel sieves and spirals of/
-- /primes/ by Colin Runciman
-- (<http://www.cs.york.ac.uk/ftpdir/pub/colin/jfp97lw.ps.gz>) and
-- /The Genuine Sieve of Eratosthenes/ by Melissa O'Neil
-- (<http://www.cs.hmc.edu/~oneill/papers/Sieve-JFP.pdf>).
-- 
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 :: Integral int => [int]
primes :: forall int. Integral int => [int]
primes = Int -> [int]
forall int. Integral int => Int -> [int]
wheelSieve Int
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
-- where @n@ is the argument given to @wheelSieve@. Don't use too
-- large wheels. The number @6@ is a good value to pass to this
-- function. Larger wheels improve the run time at the cost of higher
-- memory requirements.
-- 
wheelSieve :: Integral int
           => Int    -- ^ number of primes canceled by the wheel
           -> [int]  -- ^ infinite list of primes
wheelSieve :: forall int. Integral int => Int -> [int]
wheelSieve Int
k = [int] -> [int]
forall a. [a] -> [a]
reverse [int]
ps [int] -> [int] -> [int]
forall a. [a] -> [a] -> [a]
++ ([int] -> int) -> [[int]] -> [int]
forall a b. (a -> b) -> [a] -> [b]
map [int] -> int
forall a. HasCallStack => [a] -> a
head (int -> [int] -> [[int]]
forall int. (Ord int, Num int) => int -> [int] -> [[int]]
sieve int
p ([int] -> [int]
forall a. HasCallStack => [a] -> [a]
cycle [int]
ns))
 where (int
p:[int]
ps,[int]
ns) = Int -> ([int], [int])
forall int. Integral int => Int -> Wheel int
wheel Int
k

{-# SPECIALISE wheelSieve :: Int -> [Int]     #-}
{-# SPECIALISE wheelSieve :: Int -> [Integer] #-}

-- |
-- Checks whether a given 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 :: forall int. Integral int => int -> Bool
isPrime int
n | int
n int -> int -> Bool
forall a. Ord a => a -> a -> Bool
> int
1     = int -> [int]
forall int. Integral int => int -> [int]
primeFactors int
n [int] -> [int] -> Bool
forall a. Eq a => a -> a -> Bool
== [int
n]
          | Bool
otherwise = Bool
False

{-# 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 :: forall int. Integral int => int -> [int]
primeFactors int
n = int -> [int] -> [int]
forall {a}. Integral a => a -> [a] -> [a]
factors int
n (Int -> [int]
forall int. Integral int => Int -> [int]
wheelSieve Int
6)
 where
  factors :: a -> [a] -> [a]
factors a
1 [a]
_                  = []
  factors a
m (a
p:[a]
ps) | a
m a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
pa -> a -> a
forall a. Num a => a -> a -> a
*a
p   = [a
m]
                   | a
r a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0    = a
p a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a -> [a] -> [a]
factors a
q (a
pa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ps)
                   | Bool
otherwise = a -> [a] -> [a]
factors a
m [a]
ps
   where (a
q,a
r) = a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
quotRem a
m a
p

{-# SPECIALISE primeFactors :: Int     -> [Int]     #-}
{-# SPECIALISE primeFactors :: Integer -> [Integer] #-}

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

-- Sieves prime candidates by computing composites from the result of
-- a recursive call with identical arguments. We could use sharing
-- instead of a recursive call with identical arguments but that would
-- lead to much higher memory requirements. The results of the
-- different calls are consumed at different speeds and we want to
-- avoid multiple far apart pointers into the result list to avoid
-- retaining everything in between.
--
-- Each list in the result starts with a prime. To obtain composites
-- that need to be cancelled, one can multiply all elements of the
-- list with its head.
-- 
sieve :: (Ord int, Num int) => int -> [int] -> [[int]]
sieve :: forall int. (Ord int, Num int) => int -> [int] -> [[int]]
sieve int
p ns :: [int]
ns@(int
m:[int]
ms) = int -> [int] -> [int]
forall int. Num int => int -> [int] -> [int]
spin int
p [int]
ns [int] -> [[int]] -> [[int]]
forall a. a -> [a] -> [a]
: int -> [int] -> Composites int -> [[int]]
forall int.
(Ord int, Num int) =>
int -> [int] -> Composites int -> [[int]]
sieveComps (int
pint -> int -> int
forall a. Num a => a -> a -> a
+int
m) [int]
ms (int -> [int] -> Composites int
forall int. (Ord int, Num int) => int -> [int] -> Composites int
composites int
p [int]
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 int = (Queue int,[[int]])

-- The feeder is computed from the result of a call to 'sieve'.
-- 
composites :: (Ord int, Num int) => int -> [int] -> Composites int
composites :: forall int. (Ord int, Num int) => int -> [int] -> Composites int
composites int
p [int]
ns = (Queue int
forall int. Queue int
Empty, ([int] -> [int]) -> [[int]] -> [[int]]
forall a b. (a -> b) -> [a] -> [b]
map [int] -> [int]
forall {b}. Num b => [b] -> [b]
comps (int -> [int] -> [int]
forall int. Num int => int -> [int] -> [int]
spin int
p [int]
ns [int] -> [[int]] -> [[int]]
forall a. a -> [a] -> [a]
: int -> [int] -> [[int]]
forall int. (Ord int, Num int) => int -> [int] -> [[int]]
sieve int
p [int]
ns))
 where comps :: [b] -> [b]
comps xs :: [b]
xs@(b
x:[b]
_) = (b -> b) -> [b] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (b
xb -> b -> b
forall a. Num a => a -> a -> a
*) [b]
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 :: Ord int => Composites int -> (int,Composites int)
splitComposites :: forall int. Ord int => Composites int -> (int, Composites int)
splitComposites (Queue int
Empty, [int]
xs:[[int]]
xss) = (Queue int, [[int]]) -> (int, (Queue int, [[int]]))
forall int. Ord int => Composites int -> (int, Composites int)
splitComposites ([int] -> [Queue int] -> Queue int
forall int. [int] -> [Queue int] -> Queue int
Fork [int]
xs [], [[int]]
xss)
splitComposites (Queue int
queue, xss :: [[int]]
xss@((int
x:[int]
xs):[[int]]
yss))
  | int
x int -> int -> Bool
forall a. Ord a => a -> a -> Bool
< int
z     = (int
x, int -> (Queue int, [[int]]) -> (Queue int, [[int]])
forall int. Ord int => int -> Composites int -> Composites int
discard int
x ([int] -> Queue int -> Queue int
forall int. Ord int => [int] -> Queue int -> Queue int
enqueue [int]
xs Queue int
queue, [[int]]
yss))
  | Bool
otherwise = (int
z, int -> (Queue int, [[int]]) -> (Queue int, [[int]])
forall int. Ord int => int -> Composites int -> Composites int
discard int
z ([int] -> Queue int -> Queue int
forall int. Ord int => [int] -> Queue int -> Queue int
enqueue [int]
zs Queue int
queue', [[int]]
xss))
 where (int
z:[int]
zs,Queue int
queue') = Queue int -> ([int], Queue int)
forall int. Ord int => Queue int -> ([int], Queue int)
dequeue Queue int
queue

{-# SPECIALISE splitComposites :: Composites Int -> (Int,Composites Int) #-}
{-# SPECIALISE
    splitComposites :: Composites Integer -> (Integer,Composites Integer) #-}

-- Drops all occurrences of the given element.
--
discard :: Ord int => int -> Composites int -> Composites int
discard :: forall int. Ord int => int -> Composites int -> Composites int
discard int
n Composites int
ns | int
n int -> int -> Bool
forall a. Eq a => a -> a -> Bool
== int
m    = int -> Composites int -> Composites int
forall int. Ord int => int -> Composites int -> Composites int
discard int
n Composites int
ms
             | Bool
otherwise = Composites int
ns
 where (int
m,Composites int
ms) = Composites int -> (int, Composites int)
forall int. Ord int => Composites int -> (int, Composites int)
splitComposites Composites int
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 :: (Ord int, Num int) => int -> [int] -> Composites int -> [[int]]
sieveComps :: forall int.
(Ord int, Num int) =>
int -> [int] -> Composites int -> [[int]]
sieveComps int
cand ns :: [int]
ns@(int
m:[int]
ms) Composites int
xs
  | int
cand int -> int -> Bool
forall a. Eq a => a -> a -> Bool
== int
comp = int -> [int] -> Composites int -> [[int]]
forall int.
(Ord int, Num int) =>
int -> [int] -> Composites int -> [[int]]
sieveComps (int
candint -> int -> int
forall a. Num a => a -> a -> a
+int
m) [int]
ms Composites int
ys
  | int
cand int -> int -> Bool
forall a. Ord a => a -> a -> Bool
<  int
comp = int -> [int] -> [int]
forall int. Num int => int -> [int] -> [int]
spin int
cand [int]
ns [int] -> [[int]] -> [[int]]
forall a. a -> [a] -> [a]
: int -> [int] -> Composites int -> [[int]]
forall int.
(Ord int, Num int) =>
int -> [int] -> Composites int -> [[int]]
sieveComps (int
candint -> int -> int
forall a. Num a => a -> a -> a
+int
m) [int]
ms Composites int
xs
  | Bool
otherwise    = int -> [int] -> Composites int -> [[int]]
forall int.
(Ord int, Num int) =>
int -> [int] -> Composites int -> [[int]]
sieveComps int
cand [int]
ns Composites int
ys
 where (int
comp,Composites int
ys) = Composites int -> (int, Composites int)
forall int. Ord int => Composites int -> (int, Composites int)
splitComposites Composites int
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 :: Num int => int -> [int] -> [int]
spin :: forall int. Num int => int -> [int] -> [int]
spin int
x (int
y:[int]
ys) = int
x int -> [int] -> [int]
forall a. a -> [a] -> [a]
: int -> [int] -> [int]
forall int. Num int => int -> [int] -> [int]
spin (int
xint -> int -> int
forall a. Num a => a -> a -> a
+int
y) [int]
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 int = ([int],[int])

-- 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 :: Integral int => Int -> Wheel int
wheel :: forall int. Integral int => Int -> Wheel int
wheel Int
n = (Wheel int -> Wheel int) -> Wheel int -> [Wheel int]
forall a. (a -> a) -> a -> [a]
iterate Wheel int -> Wheel int
forall int. Integral int => Wheel int -> Wheel int
next ([int
2],[int
1]) [Wheel int] -> Int -> Wheel int
forall a. HasCallStack => [a] -> Int -> a
!! Int
n

{-# SPECIALISE wheel :: Int -> Wheel Int     #-}
{-# SPECIALISE wheel :: Int -> Wheel Integer #-}

next :: Integral int => Wheel int -> Wheel int
next :: forall int. Integral int => Wheel int -> Wheel int
next (ps :: [int]
ps@(int
p:[int]
_),[int]
xs) = (int
pyint -> [int] -> [int]
forall a. a -> [a] -> [a]
:[int]
ps,int -> int -> int -> [int] -> [int]
forall int. Integral int => int -> int -> int -> [int] -> [int]
cancel ([int] -> int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [int]
ps) int
p int
py [int]
ys)
 where (int
y:[int]
ys) = [int] -> [int]
forall a. HasCallStack => [a] -> [a]
cycle [int]
xs
       py :: int
py = int
p int -> int -> int
forall a. Num a => a -> a -> a
+ int
y

{-# SPECIALISE next :: Wheel Int     -> Wheel Int     #-}
{-# SPECIALISE next :: Wheel Integer -> Wheel Integer #-}

cancel :: Integral int => int -> int -> int -> [int] -> [int]
cancel :: forall int. Integral int => int -> int -> int -> [int] -> [int]
cancel int
0 int
_ int
_ [int]
_ = []
cancel int
m int
p int
n (int
x:ys :: [int]
ys@(int
y:[int]
zs))
  | int
nx int -> int -> int
forall a. Integral a => a -> a -> a
`mod` int
p int -> int -> Bool
forall a. Ord a => a -> a -> Bool
> int
0 = int
x int -> [int] -> [int]
forall a. a -> [a] -> [a]
: int -> int -> int -> [int] -> [int]
forall int. Integral int => int -> int -> int -> [int] -> [int]
cancel (int
mint -> int -> int
forall a. Num a => a -> a -> a
-int
x) int
p int
nx [int]
ys
  | Bool
otherwise      = int -> int -> int -> [int] -> [int]
forall int. Integral int => int -> int -> int -> [int] -> [int]
cancel int
m int
p int
n (int
xint -> int -> int
forall a. Num a => a -> a -> a
+int
yint -> [int] -> [int]
forall a. a -> [a] -> [a]
:[int]
zs)
 where nx :: int
nx = int
n int -> int -> int
forall a. Num a => a -> a -> a
+ int
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 int = Empty | Fork [int] [Queue int]

enqueue :: Ord int => [int] -> Queue int -> Queue int
enqueue :: forall int. Ord int => [int] -> Queue int -> Queue int
enqueue [int]
ns = Queue int -> Queue int -> Queue int
forall int. Ord int => Queue int -> Queue int -> Queue int
merge ([int] -> [Queue int] -> Queue int
forall int. [int] -> [Queue int] -> Queue int
Fork [int]
ns [])

{-# SPECIALISE enqueue :: [Int]     -> Queue Int     -> Queue Int     #-}
{-# SPECIALISE enqueue :: [Integer] -> Queue Integer -> Queue Integer #-}

merge :: Ord int => Queue int -> Queue int -> Queue int
merge :: forall int. Ord int => Queue int -> Queue int -> Queue int
merge Queue int
Empty Queue int
y                        = Queue int
y
merge Queue int
x     Queue int
Empty                    = Queue int
x
merge Queue int
x     Queue int
y     | Queue int -> int
forall {int}. Queue int -> int
prio Queue int
x int -> int -> Bool
forall a. Ord a => a -> a -> Bool
<= Queue int -> int
forall {int}. Queue int -> int
prio Queue int
y = Queue int -> Queue int -> Queue int
forall {int}. Queue int -> Queue int -> Queue int
join Queue int
x Queue int
y
                  | Bool
otherwise        = Queue int -> Queue int -> Queue int
forall {int}. Queue int -> Queue int -> Queue int
join Queue int
y Queue int
x
 where prio :: Queue int -> int
prio (Fork (int
n:[int]
_) [Queue int]
_) = int
n
       join :: Queue int -> Queue int -> Queue int
join (Fork [int]
ns [Queue int]
qs) Queue int
q = [int] -> [Queue int] -> Queue int
forall int. [int] -> [Queue int] -> Queue int
Fork [int]
ns (Queue int
qQueue int -> [Queue int] -> [Queue int]
forall a. a -> [a] -> [a]
:[Queue int]
qs)

{-# 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 :: forall int. Ord int => Queue int -> ([int], Queue int)
dequeue (Fork [int]
ns [Queue int]
qs) = ([int]
ns,[Queue int] -> Queue int
forall int. Ord int => [Queue int] -> Queue int
mergeAll [Queue int]
qs)

{-# SPECIALISE dequeue :: Queue Int     -> ([Int],     Queue Int)     #-}
{-# SPECIALISE dequeue :: Queue Integer -> ([Integer], Queue Integer) #-}

mergeAll :: Ord int => [Queue int] -> Queue int
mergeAll :: forall int. Ord int => [Queue int] -> Queue int
mergeAll []       = Queue int
forall int. Queue int
Empty
mergeAll [Queue int
x]      = Queue int
x
mergeAll (Queue int
x:Queue int
y:[Queue int]
qs) = Queue int -> Queue int -> Queue int
forall int. Ord int => Queue int -> Queue int -> Queue int
merge (Queue int -> Queue int -> Queue int
forall int. Ord int => Queue int -> Queue int -> Queue int
merge Queue int
x Queue int
y) ([Queue int] -> Queue int
forall int. Ord int => [Queue int] -> Queue int
mergeAll [Queue int]
qs)

{-# SPECIALISE mergeAll :: [Queue Int]     -> Queue Int     #-}
{-# SPECIALISE mergeAll :: [Queue Integer] -> Queue Integer #-}