{-# OPTIONS_GHC -fglasgow-exts #-}

module Concurrent where

import Debug.Trace
import Control.Monad.Writer
import Control.Monad.Trans
import Data.IORef

data Action m = Atom (m (Action m))
              | Fork (Action m) (Action m)
              | Stop

newtype ContT m a = ContT {runContT :: (a -> Action m) -> Action m}
instance Monad m => Monad (ContT m) where
    return x = ContT $ \k -> k x
    x >>= f = ContT $ \k -> runContT x $ \v -> runContT (f v) k

action :: Monad m => ContT m a -> Action m
action m = runContT m (\_ -> Stop)

atom :: Monad m => m a -> ContT m a
atom m = ContT $ \k -> Atom (do v <- m; return $ k v)

stop :: Monad m => ContT m a
stop = ContT $ \_ -> Stop

par :: Monad m => ContT m a -> ContT m a -> ContT m a
par x y = ContT $ \k -> Fork (runContT x k) (runContT y k)

fork :: Monad m => ContT m a -> ContT m ()
fork x = ContT $ \k -> Fork (action x) (k ())

instance MonadTrans ContT where
    lift = atom


step :: Monad m => [Action m] -> m ()
step [] = return ()
step (x:xs) = case x of
    Atom a -> do v <- a; step $ xs++[v]
    Fork a b -> step $ xs++[a,b]
    Stop -> step xs

run :: Monad m => ContT m a -> m ()
run x = step [action x]


instance MonadWriter w m => MonadWriter w (ContT m) where
    tell s = lift $ tell s

loop :: MonadWriter w m => w -> m ()
loop s = do tell s; loop s

ex1 :: ContT (Writer String) ()
ex1 = do
    tell "start!"
    fork $ loop "fish"
    loop "cat"

test1 = execWriter $ run ex1


type MVar a = IORef (Maybe a)

newMVar :: ContT (WriterT String (IO)) (MVar a)
newMVar = lift $ lift $ newIORef Nothing

writeMVar :: MVar a -> a -> ContT (WriterT String (IO)) ()
writeMVar x v = lift $ lift $ writeIORef x $ Just v

takeVar :: MVar a -> IO (Maybe a)
takeVar x = do 
    v <- readIORef x
    writeIORef x Nothing
    return v

readMVar :: MVar a -> ContT (WriterT String (IO)) a
readMVar x = do
    v <- lift $ lift $ takeVar x
    case v of
        Nothing -> readMVar x
        Just v' -> return v'


ex2 :: ContT (WriterT String (IO)) ()
ex2 = do
    tell "start!, "
    x <- newMVar
    fork (do tell "reading, "; v <- readMVar x; tell $ "found "++(show v)++", ")
    tell "writing..., "; tell "...2, "
    writeMVar x 2

test2 = do v <- execWriterT $ run ex2; putStrLn $ "Output Stream: "++v