{-# LANGUAGE Rank2Types #-}
module Control.Monad.ST.Lazy (
   ST, runST, fixST,

   RealWorld, stToIO,

   strictToLazyST, lazyToStrictST,

   unsafeInterleaveST,
   unsafeIOToST
   ) where


import qualified Control.Monad.ST.Strict as Strict
import Control.Monad.Trans.State (StateT(StateT), mapStateT, evalStateT, )
import Control.Monad.Trans.Class (lift, )
import Control.Monad.Fix (MonadFix, mfix, )
import Control.Applicative (Applicative(pure, (<*>)), )
import System.IO.Unsafe (unsafePerformIO, unsafeInterleaveIO, )



-- * The 'ST' Monad

{- |
The lazy state thread monad provides a weak form of sequencing:
If action A is written before B, then A will be run before B.
Generally an action is run whenever its result is needed
or when a subsequent action must be run.

This sequencing is necessary,
in order to run 'writeSTRef' actions at all
and run them in the right order and interleaving with 'readSTRef'.

I used the same idea in the @lazyio@ package.
-}
newtype ST s a = ST {unST :: StateT (State s) IO a}

data State s = State


instance Functor (ST s) where
   fmap f (ST st) = ST (fmap f st)

instance Applicative (ST s) where
   pure f = ST (pure f)
   ST f <*> ST x = ST $
      mapStateT unsafeInterleaveIO f <*>
      mapStateT unsafeInterleaveIO x

instance Monad (ST s) where
   return f = ST (return f)
   fail str = ST (fail str)
   ST x >> ST y = ST $
      mapStateT unsafeInterleaveIO x >>
      mapStateT unsafeInterleaveIO y
   ST x >>= k = ST $
      mapStateT unsafeInterleaveIO . unST . k =<<
      mapStateT unsafeInterleaveIO x


runST :: (forall s. ST s a) -> a
runST st = unsafePerformIO (evalStateT (unST st) State)

fixST :: (a -> ST s a) -> ST s a
fixST f = ST (mfix (unST . f))

instance MonadFix (ST s) where
   mfix = fixST


-- * Converting 'ST' to 'IO'

data RealWorld = RealWorld

stToIO :: ST RealWorld a -> IO a
stToIO = flip evalStateT State . unST


-- * Converting between strict and lazy 'ST'

strictToLazyST :: Strict.ST s a -> ST s a
strictToLazyST =
   unsafeIOToST . Strict.unsafeSTToIO

{-
ToDo:
Do we also have to bring the monadic result into some evaluated form?
Shall we also trigger all ST actions by a 'seq'ing on the final 'State'?
-}
lazyToStrictST :: ST s a -> Strict.ST s a
lazyToStrictST =
   Strict.unsafeIOToST . flip evalStateT State . unST


-- * Unsafe operations

unsafeInterleaveST :: ST s a -> ST s a
unsafeInterleaveST (ST st) =
   ST $ StateT $ \_ -> fmap (\x->(x,State)) (evalStateT st State)

{-
Explicitly matching the @State@ constructor
forces former actions to be run.
-}
unsafeIOToST :: IO a -> ST s a
unsafeIOToST m =
   ST $ StateT $ \State -> fmap (\x->(x,State)) m

{-
unsafeSTToIO :: ST s a -> IO a
unsafeSTToIO = unST
-}