Multiple independent ST/State monads within a Monad Transformer (RandT)...complicated wrapping/unwrapping

Question

Just learning how to get a deeper intuitive grasp of monads and transformers; a lot of things that might seem obvious are still kind of tricky to me haha.

So I have a computation that lives in the Rand monad, but inside it, there is another "sub-computation" (or multiple) that lives inside an ST monad (or State monad, for all it matters ... ST only for performance but I think State works just as well in this context).

The entire computation doesn't need to be inside the ST monad...and this sub-computation will be called multiple times with different starting states, so I don't want to coerce the entire thing into an ST (unless that's the idiomatic way).

Without randomness, the structure looks like this:

main = print mainComp

mainComp :: Int
mainComp = otherComp + (subComp 1) + (subComp 2)

subComp :: Int -> Int
subComp n = runST $ do
  -- generate state based on n
  -- ...
  replicateM_ 100 mutateState
  -- ...
  -- eventually returns an ST s Int

mutateState :: ST s ()
mutateState = -- ...

Basically things work pretty well and there is complete referential transparency in mainComp and subComp.

This is how I've so far used Rand --

main = (evalRandIO mainComp) >>= print

mainComp :: (RandomGen g) => Rand g Int
mainComp = do
  subResultA <- subComp 1
  subResultB <- subComp 2
  return $ otherComp + subResultA + subResultB

subComp :: (RandomGen g) => Int -> Rand g Int
subComp = return $ runST $ do           -- is this ok to just throw in return?
  -- generate state based on n
  -- ...
  replicateM_ 100 mutateState
  -- ...
  -- eventually returns an ST s Int (??)

mutateState :: ??
mutateState = ??

What is the type of mutateState supposed to be, if I wanted to use the random seed and the Rand monad in it? I think I might want to use a return type of RandT g (ST s) (), but how do I make that fit in with the type expected in the runST in subComp?

Answer 1

With monad transformers, you 'peel off' layers in the inverse order in which you add them. So, if you have something of type RandT g (ST s) (), you start by eliminating the RandT using evalRandT or runRandT, and only then calling runST.

Here's a simple example of combining RandT and ST:

import Data.STRef
import System.Random

import Control.Monad
import Control.Monad.Trans
import Control.Monad.ST
import Control.Monad.Random
import Control.Monad.Random.Class

stNrand :: RandT StdGen (ST s) Int
stNrand = do
    ref <- lift $ newSTRef 0
    i <- getRandomR (0,10)
    lift $ writeSTRef ref i
    lift $ readSTRef ref

main :: IO ()
main = putStrLn . show $ runST $ evalRandT stNrand (mkStdGen 77)

Edit: Here's an expanded version, now with a function runSTBelowRand that lets you embed RandT StdGen (ST s) a computations in Rand StdGen a computations. The function uses getSplit to split the seed of the global computation, and feed the new seed to the sub-computation.

{-# LANGUAGE RankNTypes #-}

import Data.STRef
import System.Random

import Control.Monad
import Control.Monad.Trans
import Control.Monad.ST
import Control.Monad.Random
import Control.Monad.Random.Class

stNrand :: RandT StdGen (ST s) Int
stNrand = do
    ref <- lift $ newSTRef 0
    i <- getRandomR (0,10)
    lift $ writeSTRef ref i
    lift $ readSTRef ref

runSTBelowRand :: (forall s. RandT StdGen (ST s) a) -> Rand StdGen a
runSTBelowRand r = do
    splittedSeed <- getSplit
    return $ runST $ evalRandT r splittedSeed

globalRand :: Rand StdGen (Int,Int)
globalRand = do
    i1 <- runSTBelowRand stNrand
    -- possibly non-ST stuff here
    i2 <- runSTBelowRand stNrand
    return (i1,i2)

main :: IO ()
main = putStrLn . show $ evalRand globalRand (mkStdGen 77)