gspr
gspr

Reputation: 11227

Monadic creation of vectors (or: can someone type annotate this for me?)

I came across the following piece of code as part of this Redddit sub-thread discussing an implementation of the Fisher-Yates shuffle:

randomIs g n = fill g 0
  where
    v = enumFromN 0 n

    fill g i = when (i < n) $ do
            let (x,g') = randomR (i, n-1) g
            G.swap v i x
            fill g' (i+1)

(I guess G refers to Data.Vector.Generic.Mutable... right?). Having never created vectors monadically before, I'm struggling to grasp this, especially with no type annotations. Doesn't v have type Data.Vector Int? How come one can pass it to G.swap then? Won't it have to be thawed first?

I might have just misunderstood Data.Vector.Generic, but if someone could clarify the above (by adding type annotations, perhaps?), I'd appreciate it.

Addendum: Here's my own attempt at adding type annotations:

import qualified Data.Vector.Unboxed as UVect
import qualified Data.Vector.Unboxed.Mutable as UMVect
import qualified System.Random as R
import Control.Monad
import Control.Monad.ST

randomPermutation :: forall a. (R.RandomGen a) => a -> Int -> UVect.Vector Int
randomPermutation g n = runST newVect
    where
      newVect :: ST s (UVect.Vector Int)
      newVect = UVect.unsafeThaw (UVect.enumFromN 0 n) >>= \v ->
                fill v 0 g >>
                UVect.unsafeFreeze v

      fill x i gen = when (i < n) $
                     let (j, gen') = R.randomR (i, n-1) gen in
                     UMVect.unsafeSwap x i j >>
                     fill x (i+1) gen'

As you can see, I'm avoiding Data.Vector.Generic to rule out the error source caused by perhaps not understanding it right. I'm also doing things in the ST monad.

In my head, the type of fill should be

UMVect.MVector (ST s (UVect.Vector Int)) Int -> Int -> a -> ST s ()

but GHC objects. Any hints? Again: It typechecks if I don't annotate fill.

Sidenote: I'd also like randomPermutation to return the updated random number generator. Thus, I'd need fill to also handle the generator's state. With my current type confusion, I don't see how to do that neatly. Any hints?

Upvotes: 3

Views: 451

Answers (1)

Rotsor
Rotsor

Reputation: 13783

The compile error is telling us:

Expected type: ST s (UMVect.MVector (ST s (UVect.Vector Int)) Int)
Actual type: ST s (UMVect.MVector (Control.Monad.Primitive.PrimState (ST s)) Int)

So, changing the type signature of fill to UMVect.MVector (PrimState (ST s)) Int -> Int -> a -> ST s () (adding import Control.Monad.Primitive too) solves the problem!

Upvotes: 2

Related Questions