Nice way to keep track of several references between functions in ST monad?

Question

I'm writing some code (a Metropolis-Hastings MCMC sampler) that will use a random number generator, and modify an array and potentially other structures based on this.

My initial idea was to use the ST monad, so that I could use ST arrays and the mersenne-random-pure64 package, keeping the PureMT generator as part of the state.

However I want to be able to split off some of the work into separate helper functions (e.g to sample a random integer in a given range, to update the array structure, and potentially more complicated things). To do this, I think I would need to pass the references to the PureMT gen and the array to all the functions, which could quickly become very ugly if I need to store more state.

My instinct is to group all of the state into a single data type that I can access anywhere, as I would using the State monad by defining a new datatype, but I don't know if that is possible with the ST monad, or the right way to go about it.

Are there any nice patterns for doing this sort of thing? I want to keep things as general as possible because I will probably need to add extra state and build more monadic code around the existing parts.

I have tried looking for examples of ST monad code but it does not seem to be covered in Real World Haskell, and the haskell wiki examples are very short and simple.

thanks!

Answer 1

My instinct is to group all of the state into a single data type that I can access anywhere, as I would using the State monad by defining a new datatype, but I don't know if that is possible with the ST monad, or the right way to go about it.

Are there any nice patterns for doing this sort of thing? I want to keep things as general as possible because I will probably need to add extra state and build more monadic code around the existing parts.

The key point to realize here is that it's completely irrelevant that you're using ST. The ST references themselves are just regular values, which you need access to in a variety of places, but you don't actually want to change them! The mutability occurs in ST, but the STRef values and whatnot are basically read-only. They're names pointing to the mutable data.

Of course, read-only access to an ambient environment is what the Reader monad is for. The ugly passing of references to all the functions is exactly what it's doing for you, but because you're already in ST, you can just bolt it on as a monad transformer. As a simple example, you can do something like this:

newtype STEnv s e a = STEnv (ReaderT e (ST s) a)
    deriving (Functor, Applicative, Monad)

runEnv :: STEnv s e a -> ST s e -> ST s a
runEnv (STEnv r) e = runReaderT r =<< e

readSTEnv :: (e -> STRef s a) -> STEnv s e a
readSTEnv f = STEnv $ lift . readSTRef . f =<< ask

writeSTEnv :: (e -> STRef s a) -> a -> STEnv s e ()
writeSTEnv f x = STEnv $ lift . flip writeSTRef x . f =<< ask

For more generality, you could abstract over the details of the reference types, and make it into a general "environment with mutable references" monad.

Answer 2

You can use the ST monad just like the IO monad, bearing in mind that you only get arrays and refs and no other IO goodies. Just like IO, you can layer a StateT over it if you want to thread some state transparently through your computation.