{-# LANGUAGE
 GADTs,
 EmptyDataDecls 
 #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Monad.Imperative.ImperativeMonad
-- Maintainer  :  Matthew Mirman <mmirman@andrew.cmu.edu>
-- Stability   :  experimental
-- Portability :  GADTs, EmptyDataDecls
-- Description :  A module for Imperative haskell code.
-- License     :  GNUv3
-- 
-----------------------------------------------------------------------------
module Control.Monad.Imperative.ImperativeMonad 
       ( modifyOp
       , if'
       , for 
       , break
       , continue
       , returnV
       , function 
       , auto
       , runImperative
       , liftOp2
       , prim
       , returnF
       , (=:)
       , (&)
       ) where

import Prelude hiding (break)
import Control.Monad.Cont
import Control.Monad.Reader
import Data.IORef

data Var
data Val
data Comp

data Control r = InFunction (r -> ContT r IO ())
               | InLoop { controlBreak::MIO r ()
                        , controlContinue::MIO r ()
                        , controlReturn:: r -> MIO r ()
                        }

returnF :: V a b b -> MIO b b
returnF v = do
  v' <- val v
  a <- ask
  case a of
    InLoop _ _ ret -> ret v'
    InFunction ret -> lift $ ret v'
  return v'

runImperative :: MIO a a -> IO a
runImperative foo = runContT (callCC $ \ret -> runReaderT foo $ InFunction ret) return

function :: MIO a a -> MIO b a
function = liftIO . runImperative

break :: MIO a ()
break = do
  a <- ask
  case a of
    InLoop br _ _ -> br
    _ -> return ()

continue :: MIO a ()
continue = do
  a <- ask
  case a of
    InLoop _ con _ -> con
    _ -> return ()

type MIO r a = ReaderT (Control r) (ContT r IO) a

data V b r a where
  R :: IORef a -> V Var r a
  L :: a -> V Val r a
  C :: MIO r (V b r a) -> V Comp r a

returnV a = returnF a >> return ()

val :: V b r a -> MIO r a
val v = case v of
  R r -> liftIO $ readIORef r
  L v -> return v
  C m -> val =<< m

(&) :: V Var r a -> V Var s a
(&) (R a) = R a

auto :: a -> MIO r (V Var r a)
auto a = do
  r <- liftIO $ newIORef a
  return $ R r

prim :: a -> V Val r a
prim a = L a

infixr 0 =:

(=:) :: V Var r a -> V b r a -> MIO r ()
(=:) (R ar) br = do
  b <- val br
  liftIO $ writeIORef ar b

for :: (MIO r irr1, V b r Bool, MIO r irr2) -> MIO r () -> MIO r ()
for (init, check, incr) body = init >> for'
  where for' = do
          do_comp <- val check
          when do_comp $ callCC $ \break_foo -> do
                         callCC $ \continue_foo -> do
                           flip withReaderT body $ \inbod ->
                             InLoop (break_foo ()) (continue_foo ()) (controlReturn inbod)
                         incr
                         for'

if' :: V b r Bool -> MIO r () -> MIO r ()
if' b m = do
  v <- val b
  when v m

modifyOp :: (a->b->a) -> V Var r a -> V k r b -> MIO r ()
modifyOp op (R ar) br = do
  b <- val br
  liftIO $ modifyIORef ar (\v -> op v b)

liftOp2 :: (t -> t' -> a) -> V b r t -> V b' r t' -> V Comp r a
liftOp2 foo ar br = C $ do
  a <- val ar
  b <- val br
  return $ prim $ foo a b