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

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

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.

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

instance (Monad m) => Monad (IOT m) where
	return x = IOT (\s -> return (s, x))
	IOT f >>= g = IOT (\s -> f s >>= \(s2, x) -> case g x of
		IOT h -> h 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 (\s -> liftM ((,) s) m)

instance MFunctor IOT where
	hoist f (IOT g) = IOT (f . g)

-- Flatten two layers into one. mmorph exports 'squash'.
_squash (IOT f) = IOT (\s -> let IOT g = f s in g s >>= return . snd)

instance MMonad IOT where
	embed f (IOT g) = _squash $ IOT (f . g)

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