CODEWITHSUNDEEP
haskellmonads
tsorn
tsorn

Reputation: 3625

Returning random generator from runST

I'm trying to use a random generator inside a runST statement and return the generator after use so it can be used elsewhere.

If I return only a vector, the code compiles, but when adding the generator to the return statement the compilation fails. If I understand the error message correctly, it says that the function passed to the monadic fold does not modify the vector within the same state, but I can't figure out why it then compiles if I omit the random generator from the return statement.

This compiles:

import           Control.Monad
import           Control.Monad.ST
import qualified Data.Vector.Unboxed         as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import           System.Random

randVector :: (RandomGen g) => Int -> g -> VU.Vector Int
randVector n g = runST $ do
  vector <- VU.unsafeThaw (VU.enumFromN 1 n)
  let step g i = do let (j,g') = randomR (1,n) g
                    VUM.swap vector i j
                    return g'
  g' <- foldM step g [1..VUM.length vector-1]
  VU.unsafeFreeze vector

But this does not:

randVector' :: (RandomGen g) => Int -> g -> (VU.Vector Int, g)
randVector' n g = runST $ do
  vector <- VU.unsafeThaw (VU.enumFromN 1 n) :: ST s (VUM.MVector s Int)
  let step g i = do let (j,g') = randomR (1,n) g
                    VUM.swap vector i j
                    return g'
  g' <- foldM step g [1..VUM.length vector-1]
  (VU.unsafeFreeze vector, g')

The return statement with the frozen vector and the random generator produces the following error:

Couldn't match expected type ST s (VU.Vector Int, g) with actual type (m0 (VU.Vector Int), g)

Upvotes: 2

Views: 53

Answers (2)

Benjamin Hodgson
Benjamin Hodgson

Reputation: 44614

Look at the type of unsafeFreeze.

unsafeFreeze :: (Unbox a, PrimMonad m) => MVector (PrimState m) a -> m (Vector a)

It's a monadic computation which returns m (Vector a). So the expression (VU.unsafeFreeze vector, g') has a type of (m (Vector Int), g). The type checker is telling you that it can't unify this type with the type of the do block, which you declared to be ST s (Vector Int, g).

In order to make these types unify, we need to distribute the m from inside the tuple to outside; then the type checker will infer m ~ ST s. You need to extract the Vector from its monadic computation, and then pack it back up with the generator.

do
    ...
    v <- VU.unsafeFreeze vector
    return (v, g')

Upvotes: 2

Zeta
Zeta

Reputation: 105876

VU.unsafeFreeze vector returns a vector in some monad (in this case ST s). Remember that every action has to be in the same monad in a do expression. Since (VU.unsafeFreeze vector, g') isn't of type ST s <something>, it does not work with runST.

Instead, bind the frozen vector and use return to return a single pair with type ST s (VU.Vector Int, g):

  g' <- ...
  v' <- VU.unsafeFreeze vector
  return (v', g')

Upvotes: 2

Related Questions