Passing scoped state in runST function

Question

I have a test function like below which uses runST to internally mutate state. I defined another function go within it which returns Int wrapped in ST as result (just playing with some ST concepts). The problem is that my type signature for the function seems to be wrong. If I comment out the function type signature, code runs fine. With type signature as in the commented code, it doesn't compile because the compiler interprets the state of the go function as different from the state in enclosing scope. I will appreciate pointers on how to define the function type signature to pass outer ST s to go function.

{-# LANGUAGE ScopedTypeVariables #-}
module Main where
import Data.Word(Word32)
import Data.Vector.Unboxed as U hiding (mapM_,create)
import Control.Monad.ST as ST
import Control.Monad.Primitive (PrimState)
import System.Random.MWC

test :: Word32 -> Int
test x = runST $ do
     gen <- initialize (U.singleton $ fromIntegral x :: U.Vector Word32) :: (forall s. ST s (Gen (PrimState (ST s))))
     let --go :: Int -> ST s Int
         go x = do
            v <- uniformR (1,x) gen
            return v
     i <- go 100
     return i

This is the compiler error I get if I uncomment type signature go :: Int -> ST s Int:

Couldn't match type `s1' with `s'
  `s1' is a rigid type variable bound by
       the type signature for go :: Int -> ST s1 Int at A.hs:12:16
  `s' is a rigid type variable bound by
      a type expected by the context: ST s Int at A.hs:10:10
Expected type: Gen (PrimState (ST s1))
  Actual type: Gen s
In the second argument of `uniformR', namely `gen'
In a stmt of a 'do' block: v <- uniformR (1, x) gen
In the expression:
  do { v <- uniformR (1, x) gen;
       return v }

Answer 1

The trouble is that when you say

gen <- ... :: (forall s. ST s (Gen (PrimState (ST s))))

s is now fixed to whatever s runST provides -- i.e. we cannot consider it a type variable as your signature would have you believe[1]. When the compiler says "rigid type variable", this is what it means. To reinforce that it is fixed, let's refer to it as S1 for this answer.

Notice that

let go :: Int -> ST s Int

is equivalent to

let go :: forall s. Int -> ST s Int

i.e., go must work for any s. But then you say

v <- uniformR (1,x) gen

which attempts to bind a computation of type ST S1 <something>. go is supposed to work with any s, not just S1, so this is an error. The correct signature for go is Int -> ST S1 Int, but of course we just made up S1 for argument's sake, and the true S1 has no name in the source file, so go cannot be given a signature, even though it is well-typed.

[1] Oh, you have ScopedTypeVariables on, so it looks like the forall is there because you are trying to scope s. That doesn't work -- scoped variables only apply to the body of the function with the forall. You can solve this by moving the signature to the left of the <-:

(gen :: Gen (PrimState (ST s))) <- initialize ...

after which s will be properly scoped.

Answer 2

Not a complete answer, but I can make things work by also passing gen:

test :: Word32 -> Int
test x = runST $ do
     gen <- initialize (U.singleton $ fromIntegral x :: U.Vector Word32) :: (forall s. ST s (Gen (PrimState (ST s))))
     let go3 ::  (Num b, Variate b) => Gen (PrimState (ST s)) -> b -> ST s b
         go3 g x' = uniformR (1,x') g
     go3 gen 100