{-# OPTIONS_GHC -O2 #-}
module Main where
import Data.List hiding (union)
import Data.Array.Unboxed
 
primesSA :: [Int]
primesSA = 2 : prs ()
  where 
    prs () = 3 : sieve (prs ()) 3 []
    sieve (p:ps) x fs = [i*2 + x | (i,True) <- assocs a] 
                        ++ sieve ps (p*p) ((p,0) : 
                             [(s, rem (y-q) s) | (s,y) <- fs])
     where
      q = (p*p-x)`div`2
      a :: UArray Int Bool
      a = accumArray (\ b c -> False) True (1,q-1)
                     [(i,()) | (s,y) <- fs, i <- [y+s, y+s+s..q]]

main = print $ 
         --  bprimes !!   400000 -- 100k:1.04-200k:2.69=n^1.37  -- ghc 7.6.3 !!
                                 --           400k:7.35=n^1.45  
         --  primes  !!  1000000 -- 200k:0.32-400k:0.83=n^1.38  -- CONSTANT
                                 --           1mln:2.51=n^1.21
                                 --           2mln:5.72=n^1.19  --   MEMORY !!!
            primesSA !! 10000000 -- 400k:0.43-1mln:1.12=n^1.04    
                                 --           2mln:2.26=n^1.01      9.3M
                                 --           4mln:5.10=n^1.17(1.09)    9.4M
                                 --          10mn:13.01=n^1.02(1.09,1.07)  9.4M
bprimes :: [Int]  -- Richard Bird's sieve
bprimes = _Y $ (2:) . minus [3..] . 
              foldr (\p r-> p*p : union [p*p+p, p*p+2*p..] r) []

primes :: [Int]
primes = [2,3,5,7] ++ _Y ((11:) . tail . minus (scanl (+) 11 wh11) 
               . foldi (\(x:xs) r -> x : union xs r)
               . map (\(w,p)-> scanl (\c d-> c + p*d) (p*p) w)
               . equalsBy snd (tails wh11 `zip` scanl (+) 11 wh11))

wh3  = 2:wh3                 --  ([3],2)                {1*2,2*3}
wh5  = 2:4:wh5               --  ([5,7],6)                  {2*4,6*5}
wh7  = 4:2:4:2:4:6:2:6:wh7   --  ([7,11,13,17,19,23,29,31],30)  {8*6,30*7}
wh11 = 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:wh11

_Y g = g (_Y g)        -- multistage with non-sharing fixpoint combinator
     -- = g (fix g)    -- two stages with sharing fixpoint combinator

foldi f (h:t) = f h . foldi f . unfoldr (\(a:b:c)->Just(f a b,c)) $ t

union      a b = ordzipBy id  (:)  (:)  (:)  a b
minus      a b = ordzipBy id  (:) skip skip  a b
equalsBy k a b = ordzipBy k   skip (:) skip  a b
skip       a b = b                       -- skip a = [] ; emit a = [a]

ordzipBy k f g h a b = loop a b where    -- concat$unfoldr pull(a,b)
  loop a@(x:t) b@(y:r) = case compare (k x) y of
    LT -> f x (loop t b)                 -- Just(f x,(t,b))
    EQ -> g x (loop t r)                 -- Just(g x,(t,r))
    GT -> h y (loop a r)                 -- Just(h y,(a,r))