{-# OPTIONS_GHC -O2 #-}

{-        original          tweaked          nested-feed         array-based
          (3*p,p)           (p*p,2*p)           JBwoVL             abPSOx
          6Uv0cL          2x speed-up        3x+ speed-up
                n^                n^                  n^                  n^
100K:  0.78s             0.38s               0.13s              0.065s    
200K:  2.02s   1.37      0.97s   1.35        0.29s   1.16       0.13s    1.00
400K:  5.05s   1.32      2.40s   1.31        0.70s   1.27       0.29s    1.16
800K: 12.37s   1.29                     1M:  2.10s   1.20       0.82s    1.13
 2M:                                                            1.71s    1.06
 4M:                                                            3.72s    1.12
10M:                                                            9.84s    1.06 
    overall in the tested range:
               1.33                                  1.21                1.09
    -}

module Main where   -- initial code from from stackoverflow.com/a/42122559
                    -- by stackoverflow.com/users/2144669/david-eisenstat
type Nat = Int   

data Heap = Leaf !Nat !Nat
          | Branch !Nat !Heap !Heap
          deriving Show

top :: Heap -> Nat
top (Leaf n _) = n
top (Branch n _ _) = n

leaf :: Nat -> Heap           -- tweak by stackoverflow.com/users/849891/will-ness:
leaf p = Leaf (p*p + 2*p) (2*p)        -- wn
-- leaf p = Leaf (3 * p) p             -- (original)

branch :: Heap -> Heap -> Heap
branch h1 h2 = Branch (min (top h1) (top h2)) h1 h2

pop :: Heap -> Heap           -- popAndReinsert, really
pop (Leaf n d) = Leaf (n + d) d          -- wn
-- pop (Leaf n d) = Leaf (n + 2*d) d     -- (original)
pop (Branch _ h1 h2)
  = case compare (top h1) (top h2) of
        LT -> branch (pop h1) h2
        EQ -> branch (pop h1) (pop h2)
        GT -> branch h1       (pop h2)

push :: Nat -> Heap -> Heap
push p h@(Leaf _ _)     = branch (leaf p)    h
push p (Branch _ h1 h2) = branch (push p h2) h1

primes0 :: [Nat]              -- original definition
primes0
  = let helper n h
          = case compare n (top h) of
                LT -> n : helper (n + 2) (push n h)
                EQ ->     helper (n + 2) (pop h)
                GT ->     helper  n      (pop h)
      in 2 : 3 : helper 5 (leaf 3)

primes :: [Nat]   -- nested primes feed chain by stackoverflow.com/users/849891/will-ness
primes            -- for additional 3x+ speedup due to quadratic reduction in memory usage
  = let           -- and noticeable improvement in empirical orders of growth
        helper n h q ps@(p:r)
          | n==q      = helper (n + 2) (push p h) (head r^2) r
          | otherwise = case compare n (top h) of
                LT -> n : helper (n + 2) h       q ps
                EQ ->     helper (n + 2) (pop h) q ps
                GT ->     helper  n      (pop h) q ps
        -- ps = [5,7] ++ helper 11 (leaf 3) 25 ps 
        _Y g = g (_Y g)      -- = g (g (g (g (...)))) = g . g . g . g . ....
      in -- 2 : 3 : ps 
         2 : 3 : _Y (([5,7] ++) . helper 11 (leaf 3) 25)
 
main = do (print $ primes !! 1000000)