Data.Vector.modify with nested vectors

Question

I have a vector nested inside another. I want to use modify to update this matrix in place. So I use it for the inner vector, but do I also need to use it for the outer?

Answer 1

My suggestion from the comments still stands, if you do not need to operate on a ragged array then the usual rectangular array implementation is better. Here is a short list of drawbacks of vector of vectors:

• performance penalty: the outer vector has to be boxed (which means an extra pointer indirection)
• safety: you can't guarantee the same length of all rows
• operating on ragged arrays is cumbersome

Nevertheless question still stands: how would you modify a vector of vectors in place. Below I'll provide an example function, which uses mutation to reverse rows of a ragged array and another function that reverses both rows and columns. Difference is that in the former we only mutate elements of each row, while in the latter we also mutate the outer boxed vector that corresponds to rows themselves:

{-# LANGUAGE RankNTypes #-}
import Control.Monad as M
import Control.Monad.ST
import Prelude as P
import Data.Vector as V
import Data.Vector.Generic.Mutable as VGM
import Data.Vector.Mutable as VM
import Data.Vector.Primitive as VP
import Data.Vector.Primitive.Mutable as VPM

raggedModifyRows ::
     VP.Prim a
  => (forall s. V.Vector (VPM.MVector s a) -> ST s ())
  -> V.Vector (VP.Vector a)
  -> V.Vector (VP.Vector a)
raggedModifyRows action arr = runST $ do
  -- thaw will create a copy of each row, so they can be safely modified
  mvs <- V.mapM VP.thaw arr
  action mvs
  -- We are freezing mutated copies, so it is safe to use unsafeFreeze here too
  V.mapM VP.unsafeFreeze mvs

raggedModify ::
     VP.Prim a
  => (forall s. VM.MVector s (VPM.MVector s a) -> ST s ())
  -> V.Vector (VP.Vector a)
  -> V.Vector (VP.Vector a)
raggedModify action arr = runST $ do
  arr' <- V.mapM VP.thaw arr
  -- mapM already created a copy of a boxed vector, so we can use unsafeThaw
  mv <- V.unsafeThaw arr'
  action mv
  v <- V.unsafeFreeze mv
  V.mapM VP.unsafeFreeze v

generateMatrix ::
     Prim a => (Int, Int) -> ((Int, Int) -> a) -> V.Vector (VP.Vector a)
generateMatrix (m, n) f = V.generate m $ \ i -> VP.generate n $ \j -> f (i, j)

generateRagged ::
     Prim a => V.Vector Int -> ((Int, Int) -> a) -> V.Vector (VP.Vector a)
generateRagged v f = V.imap (\ i n -> VP.generate n $ \j -> f (i, j)) v

reverseST :: (VGM.MVector v a) => v s a -> ST s ()
reverseST mv =
  let n = VGM.length mv
   in M.forM_ [0 .. (n `div` 2) - 1] $ \j -> VGM.swap mv j (n - j - 1)

reverseRaggedRows :: Prim a => V.Vector (VP.Vector a) -> V.Vector (VP.Vector a)
reverseRaggedRows = raggedModifyRows $ \rows -> V.forM_ rows reverseST

reverseRagged :: Prim a => V.Vector (VP.Vector a) -> V.Vector (VP.Vector a)
reverseRagged =
  raggedModify $ \mrows -> do
    let reverse' i = VM.read mrows i >>= reverseST
    let m = VM.length mrows
    M.forM_ [0 .. (m `div` 2) - 1] $ \i -> do
      reverse' i
      VM.swap mrows i (m - i - 1)
      reverse' i
    M.when (odd m) $ reverse' (m `div` 2)

Which can be used as follows:

λ> m = generateMatrix (3, 4) $ \(i, j) -> i+j
λ> m
[[0,1,2,3],[1,2,3,4],[2,3,4,5]]
λ> reverseRaggedRows m
[[3,2,1,0],[4,3,2,1],[5,4,3,2]]
λ> reverseRagged m
[[5,4,3,2],[4,3,2,1],[3,2,1,0]]
λ> m = generateRagged (V.fromList [1,2,3]) $ \(i, j) -> i+j
λ> m
[[0],[1,2],[2,3,4]]
λ> reverseRaggedRows m
[[0],[2,1],[4,3,2]]
λ> reverseRagged m
[[4,3,2],[2,1],[0]]

Alternatively we could have used Data.Vector.modify to operate on the outer vector or map a destructive action that uses modify across all rows. There are all sorts of ways to go about it, depends on what you are trying to achieve, for example:

λ> m = generateRagged (V.fromList [1,2,3]) $ \(i, j) -> i+j
λ> V.map (VP.modify reverseST) m
[[0],[2,1],[4,3,2]]
λ> V.modify reverseST (V.map (VP.modify reverseST) m)
[[4,3,2],[2,1],[0]]

I did recommend using massiv for regular multidimensional arrays. Therefore here is also an example of how to achieve the same with withMArrayST:

{-# LANGUAGE FlexibleContexts #-}
import Control.Monad as M
import Data.Massiv.Array as A

reverseMatrix :: Mutable r Ix2 e => Array r Ix2 e -> Array r Ix2 e
reverseMatrix arr =
  withMArrayST arr $ \marr -> do
    let Sz2 m n = msize marr
        ix2@(m2 :. n2) = m `div` 2 :. n `div` 2
    A.forM_ (0 ..: ix2) $ \ix@(i :. j) -> do
      A.swapM_ marr ix (m - i - 1 :. n - j - 1)
      A.swapM_ marr (i :. n - j - 1) (m - i - 1 :. j)
    when (odd m) $ A.forM_ (0 ..: n2) $ \ j ->
      A.swapM_ marr (m2 :. j) (m2 :. n - j - 1)
    when (odd n) $ A.forM_ (0 ..: m2) $ \ i ->
      A.swapM_ marr (i :. n2) (m - i - 1 :. n2)

Which can be used as follows:

λ> a = makeArrayR P Seq (Sz2 3 4) $ \ (i :. j) -> i + j
λ> a
Array P Seq (Sz (3 :. 4))
  [ [ 0, 1, 2, 3 ]
  , [ 1, 2, 3, 4 ]
  , [ 2, 3, 4, 5 ]
  ]
λ> reverseMatrix a
Array P Seq (Sz (3 :. 4))
  [ [ 5, 4, 3, 2 ]
  , [ 4, 3, 2, 1 ]
  , [ 3, 2, 1, 0 ]
  ]