{-# LANGUAGE MagicHash, UnboxedTuples, Rank2Types, GADTs #-}

module Control.Monad.IOT (IOT, run) where

import GHC.IO hiding (liftIO)
import GHC.Prim
import Control.Monad.Trans (MonadIO(..))
import Control.Monad.Identity
import Control.Monad.Morph
import Control.Monad
import Control.Applicative
import Unsafe.Coerce

data Ret a = Ret (State# RealWorld) a

data Sequence m where
	None :: Sequence m
	Seq :: (Monad m) => IO (Ret ()) -> Sequence (IOT m)

{-# NOINLINE runSequence #-}
runSequence :: (Monad m) => Sequence m -> State# RealWorld -> m (Ret ())
runSequence None s = return (Ret s ())
runSequence (Seq io) _ = liftIO io

-- | An IO monad transformer.
--
-- I can't run 'IOT'. Instead, I run the monad inside it.
-- This is done using 'run', and 'hoist' from mmorph.
--
-- The combination is only a monad if the parameter monad
-- isn't nondeterministic. IOT Maybe and IOT State are
-- monads, but IOT [] and IOT Cont are not.
--
-- Should be integrated with STT.

newtype IOT m t = IOT (Sequence m -> State# RealWorld -> m (Ret t))

instance (Monad m) => Monad (IOT m) where
	return x = IOT (\_ s -> return (Ret s x))
	IOT f >>= g = IOT (\i s -> f i s >>= \(Ret s2 x) -> case g x of
		IOT h -> h i s2)

instance (Monad m) => Applicative (IOT m) where
	pure = return
	(<*>) = ap

instance (Monad m) => Functor (IOT m) where
	fmap f m = m >>= return . f

instance (Monad m) => MonadIO (IOT m) where
	liftIO (IO f) = IOT (\_ s -> case f s of
		(# s2, x #) -> return (Ret s2 x))

instance MonadTrans IOT where
	lift m = IOT (\i s -> m >>= \x -> liftM (\(Ret s ()) -> Ret s x) (runSequence i s))

-- Flatten two layers into one. mmorph exports 'squash'.
--
-- Unsafely interleave actions in the outer monad, but sequence with the
-- inner monad using a sequencing fn.
_squash (IOT f) = IOT (\i s -> let IOT g = f (Seq $ IO $ \s -> (# s, Ret s () #)) s in g i s >>= \(Ret _ pr) -> return pr)

_hoist :: (forall t. m t -> n t) -> IOT m t -> IOT n t
_hoist f (IOT g) = IOT (\i s -> f (g (unsafeCoerce i) s))
-- Type safety proof: the datum i is either in None or Seq.
--   * If it is in None, it is valid at all types.
--   * If it is in Seq, the only way it can be projected is from IOT m to IO
--   and back again. liftIO is valid at both. So 'runSequence' will
--   certainly be used at a valid type.

instance MMonad IOT where
	embed f = _squash . _hoist f

instance MFunctor IOT where
	hoist = _hoist

-- | Run an IOT.
run :: IOT Identity t -> IO t
run (IOT f) = IO (\s -> case runIdentity (f None s) of
	Ret s2 x -> (# s2, x #))