#if __GLASGOW_HASKELL__ >= 700
#endif
module Math.NumberTheory.Primes.Heap (primes, sieveFrom) where
import Data.Array.Unboxed
import Data.Array.Base (unsafeAt, unsafeRead, unsafeWrite)
import Data.Array.ST
import Control.Monad.ST
import Data.List (foldl')
#if __GLASGOW_HASKELL__ < 709
import Data.Word
#endif
#ifndef SH_SIZE
#define SH_SIZE 31
#endif
data Del a = D !a !a !Int
step :: Integral a => Int -> a
step i = fromIntegral (steps `unsafeAt` i)
data Hipp a
= E
| H !a !a !Int !(Hipp a) !(Hipp a)
push :: Integral a => a -> a -> Int -> Hipp a -> Hipp a
#if __GLASGOW_HASKELL__ >= 700
push !c !p !w = go
where
less = (< c)
go (H hc hp hw l r)
| less hc = H hc hp hw (go r) l
| otherwise = H c p w (push hc hp hw r) l
go _ = H c p w E E
#else
push c p w (H hc hp hw l r)
| c < hc = H c p w (push hc hp hw r) l
| otherwise = H hc hp hw (push c p w r) l
push c p w E = H c p w E E
#endif
bubble :: Integral a => Hipp a -> Hipp a
#if __GLASGOW_HASKELL__ >= 700
bubble h@(H c p w l r) =
case r of
E -> case l of
E -> h
H lc lp lw ll lr
| lc < c -> H lc lp lw (H c p w ll lr) r
| otherwise -> h
H rc rp rw rl rr ->
case l of
H lc lp lw ll lr
| lc < c -> if lc < rc
then H lc lp lw (mkHipp c p w ll lr) r
else H rc rp rw l (mkHipp c p w rl rr)
| rc < c -> H rc rp rw l (mkHipp c p w rl rr)
| otherwise -> h
_ -> error "Heap invariant violated, left smaller than right!"
#else
bubble h@(H c p w l@(H lc lp lw ll lr) r@(H rc rp rw rl rr))
| c <= lc && c <= rc = h
| lc < rc = H lc lp lw (bubble (H c p w ll lr)) r
| otherwise = H rc rp rw l (bubble (H c p w rl rr))
bubble h@(H c p w (H lc lp lw _ _) _)
| c <= lc = h
| otherwise = H lc lp lw (H c p w E E) E
#endif
bubble h = h
#if __GLASGOW_HASKELL__ >= 700
mkHipp :: Integral a => a -> a -> Int -> Hipp a -> Hipp a -> Hipp a
mkHipp !c !p !w = go
where
less = (< c)
go l r =
case r of
E -> case l of
E -> H c p w l r
H lc lp lw _ _
| less lc -> H lc lp lw (H c p w E E) E
| otherwise -> H c p w l r
H rc rp rw rl rr ->
case l of
H lc lp lw ll lr
| less lc -> if lc < rc
then H lc lp lw (go ll lr) r
else H rc rp rw l (go rl rr)
| less rc -> H rc rp rw l (go rl rr)
| otherwise -> H c p w l r
_ -> error "Heap invariant violated, left smaller than right!"
#endif
inc :: Integral a => Hipp a -> Hipp a
inc (H c p i l r)
= bubble (H (c+p*step i) p (nextIndex i) l r)
inc h = h
adjust :: Integral a => a -> Hipp a -> Hipp a
adjust cm h@(H v _ _ _ _)
| cm == v = adjust cm (inc h)
adjust _ h = h
buildH :: Integral a => [Del a] -> Hipp a
buildH [] = E
buildH (D s p w : tl) = H s p w l r
where
(ll,rl) = goSplit [] [] tl
goSplit xs ys [] = (reverse ys, reverse xs)
goSplit xs ys (d:ds) = goSplit ys (d:xs) ds
l = buildH ll
r = buildH rl
simpleSieve :: Integral a => Hipp a -> a -> Int -> [Del a]
simpleSieve h@(H nc _ _ _ _) cd !i
| cd < nc = D s cd i : simpleSieve ( push s cd i h) (cd + step i) (nextIndex i)
| otherwise = simpleSieve (adjust cd h) (cd + step i) (nextIndex i)
where
s = cd*cd
simpleSieve _ _ _ = []
feederSieve :: Integral a => [Del a] -> Hipp a -> Hipp a -> a -> Int -> [Del a]
feederSieve dls@((D s p u):ds) sh@(H sc _ _ _ _) lh@(H lc _ _ _ _) cd i
| cd == sc = feederSieve dls (adjust cd (inc sh)) (adjust cd lh) cd' j
| cd == lc = feederSieve dls sh (adjust cd (inc lh)) cd' j
| cd == s = feederSieve ds sh (push (s + p*step u) p (nextIndex u) lh) cd' j
| otherwise = D (cd*cd) cd i : feederSieve dls sh lh cd' j
where
!cd' = cd + step i
!j = nextIndex i
feederSieve _ _ _ _ _ = []
feeder :: Integral a => a -> Int -> [Del a]
feeder p i = feederSieve lrg sh lh p i
where
(sml,D s lp w : lrg) = splitAt SH_SIZE (D q p i : simpleSieve (H q p i E E) (p+step i) (nextIndex i))
sh = buildH sml
lh = H s lp w E E
q = p*p
primeSieve :: Integral a => [Del a] -> Hipp a -> Hipp a -> a -> Int -> [a]
primeSieve dls@((D s p u):ds) sh@(H sc _ _ _ _) lh@(H lc _ _ _ _) cd i
| cd == sc = primeSieve dls ( adjust cd (inc sh)) (adjust cd lh) cd' j
| cd == lc = primeSieve dls sh (adjust cd (inc lh)) cd' j
| cd == s = primeSieve ds sh (push (s + p*step u) p (nextIndex u) lh) cd' j
| otherwise = cd : primeSieve dls sh lh cd' j
where
!cd' = cd + step i
!j = nextIndex i
primeSieve _ _ _ _ _ = []
primes :: Integral a => [a]
primes = 2:3:5:7:11:13:sieve 17 0
sieveFrom :: Integral a => a -> [a]
sieveFrom from
| fromIntegral from < (32768 :: Integer)
= dropWhile (< from) (foldr ((:) . fromIntegral) (sieve sp si) wheelPrimes)
| otherwise
= primeSieve dls sh lh start (nextIndex i0)
where
sp | odd from = 17
| otherwise = fromIntegral (remainders `unsafeAt` 0)
si | even from = 0
| otherwise = (steps `unsafeAt` 0)2
(q, r) = (from 18) `quotRem` 30030
i0 = findIx (fromIntegral r + 17)
before = 30030*q + fromIntegral (remainders `unsafeAt` i0)
!start = before + step i0
(sml, lrg) = splitAt SH_SIZE (feeder sp si)
!sh = foldl' pushD E [findMulIx p | D _ p _ <- sml]
(lh, dls) = munch E lrg
pushD h (c, p, i) = push c p i h
findMulIx p = ((p*mp), p, (nextIndex ip))
where
fpq = before `quot` p
(qq, qr) = (fpq17) `quotRem` 30030
!ip = findIx (fromIntegral qr + 17)
!mp = 30030*qq + fromIntegral (remainders `unsafeAt` ip) + step ip
munch !h dels@(D s p _ : ds)
| before < s = (h,dels)
| otherwise = munch h' ds
where
!(!c, pr, i) = findMulIx p
h' = push c pr i h
munch h [] = (h,[])
sieve :: Integral a => a -> Int -> [a]
sieve p i = primeSieve lrg sh lh p i
where
(sml,D s lp j : lrg) = splitAt SH_SIZE (feeder p i)
!sh = buildH sml
lh = H s lp j E E
nextIndex :: Int -> Int
nextIndex 5759 = 0
nextIndex i = i+1
wheelPrimes :: [Int]
wheelPrimes = 2:3:5:7:11:13:[]
findIx :: Int -> Int
findIx r
| 30030 < r = 5759
| r == m = a
| r < m = down (a1)
| otherwise = up a
where
a = min 5758 ((192*r) `quot` 1001 1)
m = remainders `unsafeAt` a
down k
| r < (remainders `unsafeAt` k) = down (k1)
| otherwise = k
up k
| r < (remainders `unsafeAt` (k+1)) = k
| otherwise = up (k+1)
remainders :: UArray Int Int
remainders = runSTUArray $ do
sar <- newArray (0,30029) True :: ST s (STUArray s Int Bool)
let n2 30030 = return ()
n2 i = unsafeWrite sar i False >> n2 (i+2)
n3 30033 = return ()
n3 i = unsafeWrite sar i False >> n3 (i+6)
n5 30035 = return ()
n5 i = unsafeWrite sar i False
>> unsafeWrite sar (i+20) False >> n5 (i+30)
n7 30037 = return ()
n7 i = unsafeWrite sar i False >> n7 (i+14)
n11 30041 = return ()
n11 i = unsafeWrite sar i False >> n11 (i+22)
n13 30043 = return ()
n13 i = unsafeWrite sar i False >> n13 (i+26)
n2 0
n3 3
n5 5
n7 7
n11 11
n13 13
rar <- newArray_ (0,5759) :: ST s (STUArray s Int Int)
let loop 30031 _ = unsafeWrite rar 5759 30031 >> return rar
loop i !r = do
c <- unsafeRead sar i
if c
then do
unsafeWrite rar r i
loop (i+2) (r+1)
else loop (i+2) r
loop 17 0
steps :: UArray Int Int
steps = runSTUArray $ do
sar <- newArray_ (0,5759) :: ST s (STUArray s Int Int)
let loop 5759 p = do
unsafeWrite sar 5759 (30047p)
return sar
loop i p = do
let !j = i+1
!n = remainders `unsafeAt` j
unsafeWrite sar i (np)
loop j n
loop 0 17