mapAccum in monad

I recently had two functions of very similar shape, only difference was that one was pure and the other need some I/O. The former was easily written using mapAccumL. I failed to find a function like mapAccumL that runs in a monad, so I wrote up the following:

mapAccumM :: (Monad m, Traversable t) => (a -> b -> m (a, c)) -> a -> t b -> m (a, t c)
mapAccumM f a l = swap <$> runStateT (mapM go l) a
        go i = do
            s <- get
            (s', r) <- lift $ f s i
            put s'
            return r

Bring on the comments/suggestions/improvements/etc!

Franklin Chen

I guess my only comment is that mapAccumL is already based on a hidden state transformer, so the generalization is natural and maybe it should even be standard?

-- left-to-right state transformer
newtype StateL s a = StateL { runStateL :: s -> (s, a) }

instance Functor (StateL s) where
    fmap f (StateL k) = StateL $ \ s -> let (s', v) = k s in (s', f v)

instance Applicative (StateL s) where
    pure x = StateL (\ s -> (s, x))
    StateL kf <*> StateL kv = StateL $ \ s ->
        let (s', f) = kf s
            (s'', v) = kv s'
        in (s'', f v)

-- |The 'mapAccumL' function behaves like a combination of 'fmap'
-- and 'foldl'; it applies a function to each element of a structure,
-- passing an accumulating parameter from left to right, and returning
-- a final value of this accumulator together with the new structure.
mapAccumL :: Traversable t => (a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL f s t = runStateL (traverse (StateL . flip f) t) s
Leave a comment