-- THE reference: Lazy Functional State Threads, PLDI 94

import GHC.Prim
import Prelude hiding (IO, init)

---------------- Control.Monad.State.Strict ----------------
                                          -- note the order is different from in ST: (value, state)
newtype State s a = State { runState :: s -> (a, s) }

instance Monad (State s) where
    return a = State $ \s -> (a, s)
    m >>= k  = State $ \s -> case runState m s of
                                 (a, s') -> runState (k a) s'

---------------- GHC.ST ----------------
-- The mental model of ST's state is a mapping from references to values.
-- newSTRef, readSTRef, and writeSTRef all operate the mapping, which is indexed by an uninstantiated type s
newtype ST s a = ST (State# s -> (# State# s, a #))

-- runState kicks off with an initial state. But runST does not take an external state parameter.
-- Intuitvely, runST always starts from an empty mapping.
-- In reality, the explicitly passed state only serves as a token for enforcing the evaluation order.
-- The implementation, in fact, directly works on a global heap.
runST :: (forall s. ST s a) -> a
runST st = case st of ST st_rep ->
                          case st_rep realWorld# of
	      		    (# _, r #) -> r

runSTBad :: ST RealWorld a -> a
runSTBad (ST st_rep) =    case st_rep realWorld# of
	      		    (# _, r #) -> r

instance Monad (ST s) where
    return x = ST $ \s -> (# s, x #)
    -- (>>=) :: ST s a -> (a -> ST s b) -> ST s b
    -- Note how the type s remains the same serving as an index
    (ST m) >>= k
      = ST $ \ s ->
	case (m s) of
          (# new_s, r #) -> case (k r) of
                              ST k2 -> k2 new_s

---------------- GHC.STRef ----------------
-- A reference is indexed by type s too, meaning it's trapped inside an ST and cannot escape from runST.
data STRef s a = STRef (MutVar# s a)

newSTRef :: a -> ST s (STRef s a)
newSTRef init = ST $ \s1# ->
    -- newMutVar# :: a -> State# s -> (# State# s,MutVar# s a #)
    -- Create MutVar# with specified initial value in specified state thread.
    case newMutVar# init s1# of
      (# s2#, var# #) -> (# s2#, STRef var# #)

---------------- GHC.IOBase ----------------
newtype IO a = IO {unIO :: State# RealWorld -> (# State# RealWorld, a #)}
-- Equivalently, type IO a = ST RealWorld a

instance  Monad IO  where
    return x	= IO $ \s -> (# s, x #)
    (IO m) >>= k     = IO ( \ s ->
                                  case m s of 
                                    (# new_s, a #) -> unIO (k a) new_s
                            )