Re: [Haskell-cafe] replicateM over vectors

1 Apr 2010

      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

Re: [Haskell-cafe] replicateM over vectors

Don Stewart