Data.Vector，与累加器映射

Question

I would like to do a map with an accumulator on Data.Vector.

I want to write the function inc:

inc :: Vector.Vector Bool -> Vector.Vector Bool

which "adds one" to the vector, eg:

inc <False, False, False> = <False, False, True>
inc <False, False, True> = <False, True, False>
inc <True, True, True> = <False, False, False>

If there was something like Data.List's mapAccumR, say with type:

mapAccumR :: (acc -> x -> (acc, y)) -> acc -> Vector x -> (acc, Vector y)

this could be done with

inc = snd . Vector.mapAccumR inc' True
  where
    inc' x y = (x && y, (x || y) && (not (x && y)))

but I can't figure it out how to do it with what is in Data.Vector.Unboxed. Is it possible?

Answer 1

The simplest solution would be to reverse your scheme and have the least significant bits on the front of the vector, like this:

inc <False, False, False> == <True, False, False>

The reason is that mapM and unfoldr are both suitable for defining inc with this bit ordering, but not with the other order, and there aren't reversed versions of these functions in vector . As an example, mapM lets us implement inc with the help of the State monad:

import Control.Monad.State
import qualified Data.Vector.Unboxed as V

inc :: V.Vector Bool -> V.Vector Bool
inc v = evalState (V.mapM go v) True where
  go acc = state $ \x -> (x /= acc, x && acc)

Alternatively, we could do two reversals to get back to the original ordering. That would be asymptotically the same, but practically significantly slower.

Of course, we can still do a lower-level implementation for mapAccumR . This necessitates working in the ST monad with mutable vectors, which isn't particularly hard, but it's not trivial either. There isn't a lot of material online on the ST monad; on Stack Overflow you may benefit from reading this question and optionally following the links from there. I try to comment on the important parts in the mapAccumR implementation below.

-- we need this so we can annotate objects in the ST monad with
-- the right parameters
{-# LANGUAGE ScopedTypeVariables #-}

import Control.Monad.ST.Strict
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Unboxed.Mutable as MV

-- note that I explicitly introduce the type variables
-- with forall. This - in conjunction with ScopedTypeVariables - 
-- lets us refer to the type variables in the function body.
mapAccumR ::
  forall x y acc.
  (V.Unbox x, V.Unbox y) =>
  (acc -> x -> (acc, y)) -> acc -> V.Vector x -> (acc, V.Vector y)
mapAccumR f acc v = runST $ do
  let len = V.length v

  -- Allocate a mutable unboxed vector of v's size.
  -- We need to annotate the "s" parameter here, so we can
  -- refer to it in the type of "go".
  (mv :: MV.STVector s y) <- MV.unsafeNew len

  -- iterate through the new vector in reverse order,
  -- updating the elements according to mapAccumR's logic.
  let go :: Int -> acc -> ST s acc
      go i acc | i < 0 = return acc
      go i acc = do
        -- y comes from the old vector
        -- we can access it via the immutable API
        let (acc' , y) = f acc (V.unsafeIndex v i)
        -- but we must do mutable writes on the new vector
        MV.unsafeWrite mv i y
        go (i - 1) acc'

  acc' <- go (len - 1) acc

  -- "unsafeFreeze" converts the mutable vector to
  -- an immutable one in-place.
  v'   <- V.unsafeFreeze mv
  return (acc', v')