{-# 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 St = St { unSt :: !(State# RealWorld) }

-- | 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.

class MFunctor t where
    {-| Lift a monad morphism from @m@ to @n@ into a monad morphism from
        @(t m)@ to @(t n)@
    -}
    hoist :: (Monad m) => (forall a . m a -> n a) -> t m b -> t n b

class (MFunctor t, MonadTrans t) => MMonad t where
    {-| Embed a newly created 'MMonad' layer within an existing layer

        'embed' is analogous to ('=<<')
    -}
    embed :: (Monad n) => (forall a . m a -> t n a) -> t m b -> t n b

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

{-#  INLINE runSequence #-}
runSequence :: (Monad m) => Sequence m -> St -> m St
runSequence None = return
runSequence (Seq io) = \_ -> liftIO io

newtype IOT m t = IOT (Sequence m -> St -> m (St, t))

instance (Monad m) => Monad (IOT m) where
	return x = IOT (\_ s -> return (s, x))
	IOT f >>= g = IOT (\i s -> f i s >>= \(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 (unSt s) of
		(# s2, x #) -> return (St s2, x))

instance MonadTrans IOT where
	lift m = IOT (\i s -> m >>= \x -> liftM (\s -> (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, St s #)) s in g i s >>= return . snd)

_hoist :: (forall t. m t -> n t) -> IOT m t -> IOT n t
_hoist f (IOT g) = IOT (\i -> f . g (unsafeCoerce i))
-- 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.
--
--   Here is the test of where things can go wrong:
test = run $ _squash $ hoist (liftIO . run) $ liftIO (print "A") >> lift (liftIO (print "B"))

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 (St s)) of
	(s2, x) -> (# unSt s2, x #))