Combining MonadRandom with ST computations in a stack

Question

I'm trying to write the Fisher-Yates shuffle using mutable arrays. So far, I have the following code:

module Main where

import Control.Monad.Random
import Control.Monad.Primitive
import Control.Monad.ST

import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV

fisherYates :: (MonadRandom m, PrimMonad m) => MV.MVector (PrimState m) a -> m ()
fisherYates v = forM_ [0 .. l - 1] (\i -> do j <- getRandomR (i, l)
                                             MV.swap v i j)
  where l = MV.length v - 1

shuffle :: MonadRandom m => V.Vector a -> m (V.Vector a)
shuffle v = _ -- don't know how to write this

main :: IO ()
main = print . evalRand (shuffle . V.generate 10 $ id) $ mkStdGen 42

However, I am totally unsure how to define shuffle, which is meant to be a 'high-level wrapper' around the mutable vector operations. It seems (at least from my understanding), that I first have to 'run' the random 'part' of the stack, save the state, run the ST 'part' to get out an immutable vector, and then rewrap the whole thing. Additionally, I know I have to make use of thaw somewhere, but my attempts are coming up short. Could someone please tell me what I'm missing, and how I can do what I want?

user_216164 · Answer

I have two suggestions for you:

Select the right monad nesting.
Separate out your monad implementation from the logic of the algorithm.

You are trying to run the random monad last and use the ST internally, thus you need the ST to be a sort of monad transformer. Decide what your monad stack looks like - which monad is the transformer and which is the inner monad? The easy thing to do is make the ST monad the inner monad and the random monad the transformer (easy just because you have all the needed packages already).

Now make a small set of helper functions. It won't really pay off here - the payoff is large for complex projects. Here's the monad stack and helpers I came up with:

{-# LANGUAGE RankNTypes #-}
module Main where

import System.Random (StdGen)
import Control.Monad.Random
import Control.Monad.Primitive
import Control.Monad.ST

import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV


type KozM s a = RandT StdGen (ST s) a

Notice the transformer is RandT and the inner monad of ST s.

rnd :: (Int,Int) -> KozM s Int
rnd = getRandomR

swp :: MV.MVector s a -> Int -> Int -> KozM s ()
swp v i j = lift (MV.swap v i j)

freeze :: MV.MVector s a -> KozM s (V.Vector    a)
thaw   :: V.Vector     a -> KozM s (MV.MVector s a)
freeze = lift . V.freeze
thaw   = lift . V.thaw

All the operations you need to mutate the vector. Now we just need to run this monad so we can somehow escape to another useful context. I noticed you previously hard-coded an RNG (42) - I use a random one but whichever...

run :: (forall s. KozM s a) -> IO a -- Can be just `a` if you hard-code
                                    -- an RNG as done in the question
run m = do g <- newStdGen
           pure (runST (evalRandT m g))

Finally we can use this monad to implement f-y:

fisherYates :: MV.MVector s a -> KozM s ()
fisherYates v = forM_ [0 .. l - 1] (\i -> do j <- rnd (i, l)
                                             swp v i j)
  where l = MV.length v - 1

At this point you might not be feeling like you learned anything, hopefully the run function was helpful but I can see how you might feel this is getting verbose. The important thing to take note of here is how clean the rest of your code can be if you handle the plumbing of the monad above so you don't have lift and module qualifiers polluting the logic of the possibly complex thing you actually set-out to solve.

That said, here's the underwhelming shuffle:

shuffle :: V.Vector a -> KozM s (V.Vector a)
shuffle v = do
    vm <- thaw v
    fisherYates vm
    freeze vm

The type is important. You had previously called evalRand on shuffle which implied it would be some sort of MonadRandom m and simultaneously have to call runST - a conflation of the monad logic and the algorithm logic that just hurts the brain.

The main is uninteresting:

main :: IO ()
main = print =<< (run (shuffle (V.generate 10 id)) :: IO (V.Vector Int))

EDIT: yes you can do this while keeping the methods more general. At some point you need to specify which monad you run or you can't have a main that will execute it, but all the logic can be general using MonadRandom and PrimMonad.

{-# LANGUAGE RankNTypes #-}
module Main where

import System.Random (StdGen)
import Control.Monad.Random
import Control.Monad.Primitive
import Control.Monad.ST

import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
type KozM s a = RandT StdGen (ST s) a

rnd  :: MonadRandom m => (Int, Int) -> m Int
rnd = getRandomR

swp :: PrimMonad m => MV.MVector (PrimState m)  a -> Int -> Int -> m ()
swp v i j = MV.swap v i j

-- freeze :: MV.MVector s a -> KozM s (V.Vector    a)
-- thaw   :: V.Vector     a -> KozM s (MV.MVector s a)
freeze :: PrimMonad m => MV.MVector (PrimState m) a -> m (V.Vector a)
thaw :: PrimMonad m => V.Vector a -> m (MV.MVector (PrimState m) a)
freeze = V.freeze
thaw   = V.thaw


-- Some monad libraries, like monadlib, have a generalized `run` class method.
-- This doesn't exist, to the best of my knowledge, for the intersection of ST
-- and mtl.
run :: (forall s. KozM s a) -> IO a -- Can be just `a` if you hard-code
                                    -- an RNG as done in the question
run m = do g <- newStdGen
           pure (runST (evalRandT m g))

-- fisherYates :: MV.MVector s a -> KozM s ()
fisherYates :: (MonadRandom m, PrimMonad m) => MV.MVector (PrimState m) a -> m ()
fisherYates v = forM_ [0 .. l - 1] (\i -> do j <- rnd (i, l)
                                             swp v i j)
  where l = MV.length v - 1

shuffle :: (MonadRandom m, PrimMonad m) => V.Vector a -> m (V.Vector a)
shuffle v = do
    vm <- thaw v
    fisherYates vm
    freeze vm

main :: IO ()
main = print =<< (run (shuffle (V.generate 10 id)) :: IO (V.Vector Int))

Combining MonadRandom with ST computations in a stack

1 Answers