-----------------------------------------------------------------------------
-- |
-- Module      :  CorecQueues
-- Copyright   :  (c) Leon P Smith 2009
-- License     :  BSD3
--
-- Maintainer  :  leon at melding-monads dot com
-- Stability   :  experimental
-- Portability :  portable
--
-----------------------------------------------------------------------------

{-# OPTIONS_GHC -O2 #-}

{-# LANGUAGE MultiParamTypeClasses       #-}
{-# LANGUAGE FunctionalDependencies      #-}
{-# LANGUAGE FlexibleInstances           #-}
{-# LANGUAGE FlexibleContexts            #-}
{-# LANGUAGE RankNTypes                  #-}

{-# LANGUAGE NoMonomorphismRestriction   #-}
{-# LANGUAGE GeneralizedNewtypeDeriving  #-}

{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE RecursiveDo   #-}

module CorecQueues where

import Control.Monad.Cont                 -- The Monad Template Library
import Control.Monad.Reader
import Control.Monad.Writer
import Control.Monad.State

import Control.Monad.ST.Strict
import Data.STRef.Strict

import Control.Monad(mapM_, liftM)        -- A few basic library functions
import Control.Arrow((***), (&&&), (>>>)) 
import Data.List(zipWith3)
import Data.Ratio

import Data.IORef                         -- For the thought experiment
import System.IO.Unsafe
import qualified Debug.Trace

import GHC.Exts -- (inline)

trace   = Debug.Trace.trace

data  Tree a b 
   =  Leaf    a   
   |  Branch  b (Tree a b) (Tree a b) 
      deriving (Eq,Show)

fib :: Int -> Tree Int Int
fib n = fibs !! (n-1)
  where 
   fibs = Leaf 0 : Leaf 0 : zipWith3 Branch [1..] fibs (tail fibs)

sternBrocot :: Tree a (Ratio Integer)
sternBrocot = loop 0 1 1 0
   where loop a b x y 
          = Branch (m%n) (loop a b m n) (loop m n x y)
            where  m  = a + x
                   n  = b + y
bot :: forall a. a
bot = error "bottom"
labelDisj :: (a -> c) -> (b -> c) -> Tree a b -> c
labelDisj leaf branch (Leaf    a      ) = leaf    a
labelDisj leaf branch (Branch  b _ _  ) = branch  b

childrenOf :: Tree a b ->  [Tree a b]
childrenOf (Leaf    _      )  =   []
childrenOf (Branch  _ l r  )  =   [l,r]
levelOrder ::  Tree a b  ->  [Tree a b]
levelOrder     tree      =   queue 
  where
    queue = tree : explore 1 queue

    explore ::  Int ->   [ Tree a b ]                  ->  [ Tree a b ]
    explore     0                            head      =            []
    explore     (n+1)    ( Branch  _ l r  :  head'  )  =   l : r :  explore (  n+2)  head'
    explore     (n+1)    ( Leaf    _      :  head'  )  =            explore    n     head'
isBranch = labelDisj (const False) (const True)

levelOrder' ::  Tree a b  ->  [Tree a b]
levelOrder'     tree      =   queue
  where
    queue  | isBranch tree   = tree :  explore 1 queue
           | otherwise       =         explore 0 queue
 
    explore :: Int -> [Tree a b] -> [Tree a b]  
    explore  0 _ = []
    explore  (n+1) (Branch _ l r : head') 
      =  if (isBranch  l) 
         then   if (isBranch r)
                then   l :  r :  explore  (  n+2)  head'
                else   l :       explore  (  n+1)  head'
         else   if (isBranch r)
                then        r :  explore  (  n+1)  head'
                else             explore     n     head'
levelOrder'2 ::  Tree a b  ->  [  Tree a b  ]
levelOrder'2     tree      =      queue
  where
     queue = enqs [tree] 0 queue
  
     enqs  ::  [Tree a b] ->  Int ->  [ Tree a b ]     ->  [ Tree a b ]
     enqs      []             n              head      =        deq          n      head
     enqs      (t:ts)         n              head
               | isBranch t                            =   t :  enqs  ts  (  n+1)   head
               | otherwise                             =        enqs  ts     n      head
 
     deq                      0              _         =        []
     deq                      (n+1)   ( t :  head'  )  =        enqs  ts'    n      head'
               where ts' = childrenOf t
prop_levelOrder'2 n tree = 
        take n ( map (labelDisj id id) (levelOrder'   tree) )
    ==  take n ( map (labelDisj id id) (levelOrder'2  tree) )
class Queue q where
   empty :: q e
   enque :: e -> q e -> q e
   deque :: q e -> (Maybe e, q e)

class Monad m => MonadQueue e m |  m -> e where
   enQ  :: e -> m ()
   deQ  :: m (Maybe e)
levelOrder'3 t = q
  where
    q = enq1 [t] (0::Int) q

    enq2 [a,b]  n head  | isBranch a  = a :  (enq1 [b]  )  (  n+1)   head
                        | otherwise   =      (enq1 [b]  )     n      head

    enq1 [a]    n head  | isBranch a  = a :  (deq       )  (  n+1)   head
                        | otherwise   =      (deq       )     n      head

    deq       0                      _       = []   
    deq       (n+1) (Branch _ l r :  head')  = enq2 [l,r] n head'

levelOrder'4 t = q
  where
    q = (enq t $ deq) (0::Int) q

    enq a k n q | isBranch a  = a :  k  (  n+1  )  q
                | otherwise   =      k     n       q

    deq    0           _                    =  []
    deq (  n+1  )   (  Branch _ l r : q  )  =  (enq l $ enq r $ deq ) n q
prop_levelOrder'3 n tree =
        take n ( map (labelDisj id id) (levelOrder'   tree) ) 
    ==  take n ( map (labelDisj id id) (levelOrder'3  tree) )

prop_levelOrder'4 n tree =
        take n ( map (labelDisj id id) (levelOrder'   tree) ) 
    ==  take n ( map (labelDisj id id) (levelOrder'4  tree) )

prop_levelOrder'5 tree n =
        take n ( map (labelDisj id id) (levelOrder'   tree) )
    ==  take n ( map (labelDisj id id) (levelOrder'5  tree) )
levelOrder'5 t = q
  where
    q = (handle t (\() -> explore)) (0::Int) q 

    handle t  | isBranch t   = enq t
              | otherwise    = ret ()
    
    explore =  deq        (\(Branch _ l r) -> 
               handle l   (\() -> 
               handle r   (\() -> 
               explore    )))
                          
    enq e  k        n            q     =  e :  k  ()  (  n+1)  q

    ret a  k        n            q     =       k  a      n     q 
        
    deq    k        0            q     =       []
    deq    k   (n+  1  )  ( e :  q' )  =       k  e      n     q'     
type CorecQSt e = Int -> [e] -> [e]

newtype  CorecQ e a  
      =  CorecQ { unCorecQ :: (a -> CorecQSt e) -> CorecQSt e }

instance Monad (CorecQ e) where
  return a  = CorecQ (\k -> k a)
  m >>= f   = CorecQ (\k -> unCorecQ m (\a -> unCorecQ (f a) k))

instance MonadQueue e (CorecQ e) where
  enQ e = CorecQ (\k n q -> e : (k () $! n+1) q)
  deQ   = CorecQ deq
    where
      deq k      0       q    = k Nothing   0 q
      deq k (n+  1) (e:  q')  = k (Just e)  n q'

runCorecQ :: CorecQ element result -> [element]
runCorecQ m =  q
        where  q = unCorecQ m (\a n' q' -> []) 0 q
levelOrder'' :: MonadQueue (Tree a b) q => Tree a b -> q ()
levelOrder'' t = handle t >> explore
   where 
     handle t  | isBranch t  = enQ t
               | otherwise   = return ()
     explore = deQ >>=  maybe (return ()) 
                        (\(Branch _ l r)  -> 
                             do
                               handle l
                               handle r
                               explore    )
levelOrder_'' t = runCorecQ (handle t >> explore)
   where 
     handle t  | isBranch t  = enQ t
               | otherwise   = return ()
     explore = deQ >>=  maybe (return ()) 
                        (\(Branch _ l r)  -> 
                             do
                               handle l
                               handle r
                               explore    )
byLevel :: (MonadQueue a m) => (a -> [a]) -> [a] -> m ()
byLevel childrenOf as = mapM_ enQ as >> explore
  where
    explore = deQ >>=  maybe (return ()) 
                       (\a ->  do 
                               mapM_ enQ (childrenOf a  )
                               explore                  )


byLevel' :: (MonadQueue a m) => (a -> [a]) -> [a] -> m ()
byLevel' childrenOf as = mapM_ handle as >> explore
  where
    handle a = when (hasChildren a) (enQ a)

    explore = deQ >>=  maybe (return ()) 
                       (\a ->  do 
                               mapM_ handle (childrenOf a  )
                               explore                     )

    hasChildren = not . null . childrenOf
data   List     st a = Null | Cons a !(ListPtr st a)
type   ListPtr  st a = STRef st (List st a)

type     STQSt st r e = ListPtr st e -> ListPtr st e -> ST st r
newtype  STQ e a 
      =  STQ { unSTQ :: forall r st. ((a -> STQSt st r e) -> STQSt st r e) }

instance Monad (STQ e) where
    return a  = STQ (\k -> k a)
    m >>= f   = STQ (\k -> unSTQ m (\a -> unSTQ (f a) k))

instance MonadQueue e (STQ e) where
    enQ e  = STQ $ \k  a z -> do
                       z' <- newSTRef Null
                       writeSTRef z (Cons e z')
                       k () a z'

    deQ    = STQ $ \k  a z -> do
                       list <- readSTRef a
                       case list of
                         Null         -> k Nothing  a  z
                         (Cons e a')  -> k (Just e) a' z

runSTQ :: STQ element result -> result 
runSTQ m = runST $ do
  ref <- newSTRef Null
  unSTQ m (\r _a _z -> return r) ref ref
foldrByLevel  :: (MonadQueue a m) 
              => (a -> [a]) -> (a -> b -> b) -> b -> [a] -> m b
foldrByLevel  childrenOf f b as = mapM_ enQ as >> explore
   where
     explore = deQ >>=  maybe (return b)
                        (\a ->  do
                                mapM_ enQ (childrenOf a  )
                                b <- explore
                                return (f a b)           )

prop_foldrByLevel childrenOf f b as =
       foldr f b (runCorecQ (byLevel childrenOf as))
   ==  runSTQ (foldrByLevel childrenOf f b as)
cid = const id

getUnion     f = f childrenOf (labelDisj  (:)  (:)  ) []
getLeaves    f = f childrenOf (labelDisj  (:)  cid  ) []
getBranches  f = f childrenOf (labelDisj  cid  (:)  ) []
foldrByLevel'  :: (MonadQueue a m) 
               => (a -> [a]) -> (a -> b -> b) -> b -> [a] -> m b 
foldrByLevel'  childrenOf f b as = handleMany as
   where
     handleMany []      = explore
     handleMany (a:as)  = do
         when (hasChildren a) (enQ a) 
         b <- handleMany as 
         return (f a b)

     explore = deQ >>= maybe (return b) (handleMany . childrenOf) 

     hasChildren = not . null . childrenOf

prop_foldrByLevel' childrenOf f b as =
       runSTQ (foldrByLevel   childrenOf f b as)
   ==  runSTQ (foldrByLevel'  childrenOf f b as)
data TwoStackQ e = TwoStackQ [e] [e]

instance Queue TwoStackQ where 
   empty = TwoStackQ [] []

   enque z  (TwoStackQ []     [])  =   TwoStackQ [z]     []
   enque z  (TwoStackQ (a:as) zs)  =   TwoStackQ (a:as)  (z:zs)

   deque    (TwoStackQ []     [])  =   ( Nothing  ,  TwoStackQ []   []  )
   deque    (TwoStackQ (a:as) zs)  
            |  null as             =   ( Just a   ,  TwoStackQ as'  []  )
            |  otherwise           =   ( Just a   ,  TwoStackQ as   zs  )
               where as' = reverse zs
newtype StateQ e a = StateQ (State (TwoStackQ e) a) deriving (Monad)

instance MonadQueue e (StateQ e) where
   enQ = StateQ . modify . enque
   deQ = StateQ (State deque)


runStateQ :: StateQ element result -> result
runStateQ (StateQ m) = let  (result, finalQ) = runState m empty
                        in  result
newtype  DebugQ q e a 
      =  DebugQ (StateT (q e) (Writer [e])  a)  
         deriving (Monad)

instance Queue q => MonadQueue e (DebugQ q e) where 
    enQ e = DebugQ (tell [e] >> modify (enque e))
    deQ   = DebugQ (StateT (return . deque))

runDebugQ :: (Queue q) => DebugQ q element result -> (result, [element])
runDebugQ (DebugQ m) 
        = let ((result, final_queue), queue) = runWriter (runStateT m empty)
           in (result, queue)
newtype CpSt s a 
      = CpSt { unCpSt :: forall r. (a -> s -> r) -> s -> r }

instance Monad (CpSt s) where
   return a  = CpSt (\k -> k a)
   m >>= f   = CpSt (\k -> unCpSt m (\a -> unCpSt (f a) k))

instance MonadState s (CpSt s) where
   get     = CpSt (\k s -> k s   s   )
   put s'  = CpSt (\k _ -> k ()  s'  )

runCpSt :: CpSt s a -> s -> (a,s)
runCpSt m s0 = unCpSt m (\a s -> (a,s)) s0
newtype CpStQ r e a = CpStQ (StateT (TwoStackQ e) (Cont r) a) deriving (Monad)

instance MonadQueue e (CpStQ r e) where
   enQ = CpStQ . modify . enque
   deQ = CpStQ (StateT (return . deque))

runCpStQ :: CpStQ r e r -> r
runCpStQ (CpStQ m) = runCont (runStateT m empty) (\(r,_q) -> r)
class MonadMapCC a m | m -> a where
    mapCC :: (a -> a) -> m b -> m b

instance MonadMapCC r (CpStQ r e) where  
    mapCC f (CpStQ m) = CpStQ (StateT (\s -> mapCont f (runStateT m s)))


foldrByLevel''  :: ( MonadQueue a m
                   , MonadMapCC b m )
                => (a -> [a]) -> (a -> b -> b) -> b -> [a] -> m b 
foldrByLevel''  childrenOf f b as = handleMany as
   where
     handleMany []      = explore
     handleMany (a:as)  = do
         when (hasChildren a) (enQ a) 
         mapCC (f a) (handleMany as)

     explore = deQ >>= maybe (return b) (handleMany . childrenOf) 

     hasChildren = not . null . childrenOf
newtype CpSt' s a 
      = CpSt' { runCpSt' :: forall r. (forall r'.  a -> s -> r') -> s -> r }

mapCpSt' :: (forall a. a -> a) -> CpSt' s b -> CpSt' s b
mapCpSt' f m  = CpSt' (\k -> f (runCpSt' m k))   
getReader    = lift ask

setReader    = modReader . const

modReader f  = callCC (\k -> local f (k ()))

stepReader f = do
                 st <- getReader
                 let (r,st') = f st
                 setReader st'
                 return r
data Len a = Len !Int a

deQ_len (Len  0          q    ) = (Nothing  , Len 0 q   )
deQ_len (Len  (n+1) (e:  q')  ) = (Just e   , Len n q'  )

inc_len (Len n head) = Len (n+1) head
newtype CorecQ' e a 
      = CorecQ' { unCorecQ' :: ContT [e] (Reader (Len [e])) a }
        deriving (Monad)

instance MonadQueue e (CorecQ' e) where
    enQ e  = CorecQ' (mapContT (liftM (e:)) (modReader inc_len))
    deQ    = CorecQ' (stepReader deQ_len)

runCorecQ' (CorecQ' m) = q 
   where  q = runReader (runContT m endpoint) (Len 0 q)
          endpoint _ = return []
newtype CorecQT e m a 
      = CorecQT ( ContT (m [e]) (Reader (Len [e])) a )
        deriving (Monad)

instance Monad m => MonadQueue e (CorecQT e m) where
   enQ z = CorecQT (mapContT (liftM (liftM (z:))) (modReader inc_len))
   deQ   = CorecQT (stepReader deQ_len)

runCorecQT m  = mfix (\q -> run m end_point (Len 0 q))
  where
    end_point _ = return []
    run (CorecQT m) k st = runReader (runContT m k) st
newtype  CorecQW w e a 
      =  CorecQW { unCorecQW :: ContT ([e],w) (Reader (Len [e])) a }
         deriving (Monad)

instance MonadQueue e (CorecQW w e) where
  enQ e = CorecQW (mapContT (liftM ((e:) *** id)) (modReader inc_len))
  deQ   = CorecQW (stepReader deQ_len)

instance MonadMapCC w (CorecQW w e) where
  mapCC f = CorecQW . mapContT (liftM (id *** f)) . unCorecQW

runCorecQW :: CorecQW w e w -> ([e],w)
runCorecQW m = (q,w)
    where  (q,w) = run m end_point (Len 0 q)
           end_point w = return ([],w)
           run m k st = runReader (runContT (unCorecQW m) k) st
type QSt r e  = IORef r -> IORef [e] -> Int -> [e] -> [e]

newtype Q r e a = Q { unQ :: ((a -> QSt r e) -> QSt r e) }

instance Monad (Q r e) where
    return a = Q ($a)
    m >>= f  = Q (\k -> unQ m (\a -> unQ (f a) k))

unsafeRead  ref   = unsafePerformIO (readIORef  ref  )
unsafeWrite ref a = unsafePerformIO (writeIORef ref a)
unsafeNew   a     = unsafePerformIO (newIORef   a    )

instance Show e => MonadQueue e (Q r e) where
  enQ x = Q (\k r e !n xs ->  let  xs' = (k () r e $! n+1) xs
                              in   trace ("enQ $ " ++ show x) 
                                   (unsafeWrite e xs' `seq` (x:xs'))  )
  deQ   = Q delta
    where
      delta k r e      0        xs   = k Nothing r e 0 xs
      delta k r e (n+  1)  (x:  xs)  = trace  ("deQ " ++ show x) 
                                              (k (Just x) r e n xs)

runQ m = (trace "reading return value" `seq` unsafeRead r (), queue)
   where
     r = unsafeNew init
     init () = unsafePerformIO $ do
                 trace "forcing computation\n" (return ())
                 xs <- readIORef e
                 force xs
                 trace "reading return value\n" (return ())
                 f <- readIORef r
                 return (f ())
     e  = unsafeNew queue 
     queue = unQ m breakK r e 0 queue

force []     = return ()
force (_:_)  = return ()

breakK a r e n xs = trace  ("setting return value: " ++ show a)
                           (unsafeWrite r (\() -> a) `seq` [])