{-# OPTIONS_GHC -O2 -v -rtsopts -with-rtsopts="-sstderr" #-}

module Main where

type Nat = Int                -- stackoverflow.com/a/42122559/849891

data Heap = Leaf !Nat !Nat    -- by stackoverflow.com/users/2144669/david-eisenstat
          | Branch !Nat !Heap !Heap
          deriving Show

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

leaf :: Nat -> Heap  
leaf p = Leaf (3 * p) p

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 + 2*d) d
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)

main = do 
          print (primes0 !! 800000)