How would I write imapInto using mutable vectors?

Question

This is related to this question:

Don't know where to start with mutable Vectors

How would I write a function that took values from one vector, transformed them and put the results into a second vector? More specifically, it should iterate over all the indexes of a source array, and hand that index and the array to a function, then store the result of the function in another array.

I think the signature would look something like this:

imapInto :: (PrimMonad m, Unbox a) => (MVector (PrimState m) a -> Int -> a) -> MVector (PrimState m) a -> MVector (PrimState m) a  -> m ()

And would be called something like this:

import Data.Vector.Unboxed.Mutable MV

...
let dst = MV.replicate 10 0.0 in
let src = MV.replicate 10 1.0 in
imapInto (\src' i -> (src' * 2.0)) dst src
...

Which would multiply all the elements in src by 2 and put the results into dst, yielding a state unit.

Answer 1

I don't think there is a builtin function which do what you want. But with these primitives it should be possible to write imperative code for your task. And if it will be general enough, you should be able even use it in pure code with create. I guess it will be something like (just sketch, untested) this:

imapInto :: (PrimMonad m, Unbox a) => (MVector (PrimState m) a -> Int -> a) -> MVector (PrimState m) a -> MVector (PrimState m) a -> m ()
imapInto f d s = go 0
    where
        go i = when (i < length d) $ write d i (f s i) >> go (i+1)

Your type of mapping function seems weird though. Didn't you mean a -> Int -> a? Then the code above will require some little changes.

UPDATE

Here is the example of usage, along with updated version of the above function. The least possible tweak was to add m type constructor to mapper function:

module Main where

import qualified Data.Vector.Unboxed.Mutable as MV
import qualified Data.Vector.Unboxed as U
import Control.Monad (forM_)
import Control.Monad.Primitive

imapInto :: (PrimMonad m, MV.Unbox a) 
         => (MV.MVector (PrimState m) a -> Int -> m a)
         -> MV.MVector (PrimState m) a 
         -> MV.MVector (PrimState m) a 
         -> m ()
imapInto f d s = forM_ [0..MV.length s - 1] $ \i -> do
    v <- f s i
    MV.write d i v

main = do
    -- Create two vectors
    v1 <- MV.replicate 10 1
    v2 <- MV.new 10
    -- Map v1 into v2 using function mapper defined below
    imapInto mapper v2 v1
    -- Print the source and the result
    uv1 <- U.unsafeFreeze v1
    uv2 <- U.unsafeFreeze v2
    print $ U.toList uv1   -- [1,1,1,1,1,1,1,1,1,1]
    print $ U.toList uv2   -- [0,1,2,3,4,5,6,7,8,9]

    where
        -- Mapper reads a value from the vector and multiplies it by its index
        mapper v i = fmap (*i) $ MV.read v i

Since you're beginner I tried to be as simple as possible, ask if something is not clear.

As you can see, I took advantage of Louis' comment and used forM_ function to make imapInto even more simpler (forM_ is mapM_ with arguments swapped). It now looks like ordinary for loop from imperative languages. Also as I said I changed mapping function type from (MVector (PrimState m) a -> Int -> a) to (MVector (PrimState m) a -> Int -> m a). This is required because you cannot do much with a mutable vector (it is passed as first parameter here) outside its monad. But inside the monad we can do anything with it, and here we are simply reading i-th element and multiplying it by i.

Note that in this version of function you can do anything with source vector from inside the mapping function. If you do not require this and you only use i-th element in conjunction with i itself, you can simplify this even more:

imapInto' :: (PrimMonad m, MV.Unbox a, MV.Unbox b) 
         => (a -> Int -> b)
         -> MV.MVector (PrimState m) b
         -> MV.MVector (PrimState m) a
         -> m ()
imapInto' f d s = forM_ [0..MV.length s - 1] $ \i -> do
    v <- MV.read s i
    MV.write d i (f v i)

Here mapping function is pure and takes only i-th element as first argument, not the vector itself. Also I've generalized this version, so it can map first vector into a vector with another type, so source and destination vectors are not required to be of the same type. This is something which could be done to the previous version, too.