[Haskell-cafe] replicateM over vectors

[Don Stewart]

Chad.Scherrer:
> Hi,
> 
> I'd like to be able to do replicateM, but over a vector instead of a list. Right now I'm doing this:
>  
> import qualified Data.Vector.Generic as G
> import qualified Data.Vector.Unboxed.Mutable as M
> replicateM n action = do
>   mu <- M.unsafeNew n
>   let go !i | i < n = action >>= M.unsafeWrite mu i >> go (i+1)
>              | otherwise = G.unsafeFreeze mu
>   go 0
> 
> I thought it might be useful to express this in terms of the available
> primitives, since this might fuse more easily, but I don't yet see a
> way to do it.
> 
> Is there a better (more elegant and/or faster) way to write this?
> 

The way you're writing it is fine and direct, but yes, I think the
challenge is to have have monadic generators that will fuse.

The direct style, I'd write as:

    import qualified Data.Vector.Generic         as G
    import qualified Data.Vector.Generic.Mutable       as M
    import qualified Data.Vector.Generic.New           as N
    import qualified Data.Vector.Fusion.Stream.Monadic       as S

    replicateM  :: (G.Vector v a) => Int -> IO a -> IO (v a)
    replicateM n a = do
        v  <- M.new n
        fill v 0
        G.unsafeFreeze v
      where
        fill v i
            | i < n = do
                x <- a
                M.unsafeWrite v i x
                fill v (i+1)
            | otherwise = return ()

But that doesn't fuse, obviously. So it seems we should be able to do something like:

    -- Almost there:
    replicateS n a = {-unsafeFreeze-} (M.unstream (S.replicate n a))

Or use Data.Vector.Generic.New, but I couldn't get the monadic action to work
out, ending up in a twisty maze of PrimMonads.

Roman? Can we generate frozen arrays for monadic generators, and still fuse in
the current New/Mutable/MStream architecture?

-- Don