Skip to content

Commit

Permalink
Add mapAccumM, forAccumM to Data.Traversable
Browse files Browse the repository at this point in the history
Approved by Core Libraries Committee in
haskell/core-libraries-committee#65 (comment)
  • Loading branch information
lykahb authored and Marge Bot committed Jul 19, 2022
1 parent aa75bbd commit 4b98c5c
Show file tree
Hide file tree
Showing 3 changed files with 80 additions and 5 deletions.
35 changes: 33 additions & 2 deletions libraries/base/Data/Functor/Utils.hs
Expand Up @@ -10,8 +10,8 @@
module Data.Functor.Utils where

import Data.Coerce (Coercible, coerce)
import GHC.Base ( Applicative(..), Functor(..), Maybe(..), Monoid(..), Ord(..)
, Semigroup(..), ($), otherwise )
import GHC.Base ( Applicative(..), Functor(..), Maybe(..), Monad (..)
, Monoid(..), Ord(..), Semigroup(..), ($), liftM, otherwise )
import qualified GHC.List as List

-- We don't expose Max and Min because, as Edward Kmett pointed out to me,
Expand Down Expand Up @@ -95,6 +95,37 @@ instance Applicative (StateR s) where
(s'', x) = kx s'
in (s'', f x y)

-- | A state transformer monad parameterized by the state and inner monad.
-- The implementation is copied from the transformers package with the
-- return tuple swapped.
--
-- @since 4.18.0.0
newtype StateT s m a = StateT { runStateT :: s -> m (s, a) }

-- | @since 4.18.0.0
instance Monad m => Functor (StateT s m) where
fmap = liftM
{-# INLINE fmap #-}

-- | @since 4.18.0.0
instance Monad m => Applicative (StateT s m) where
pure a = StateT $ \ s -> return (s, a)
{-# INLINE pure #-}
StateT mf <*> StateT mx = StateT $ \ s -> do
(s', f) <- mf s
(s'', x) <- mx s'
return (s'', f x)
{-# INLINE (<*>) #-}
m *> k = m >>= \_ -> k
{-# INLINE (*>) #-}

-- | @since 4.18.0.0
instance (Monad m) => Monad (StateT s m) where
m >>= k = StateT $ \ s -> do
(s', a) <- runStateT m s
runStateT (k a) s'
{-# INLINE (>>=) #-}

-- See Note [Function coercion]
(#.) :: Coercible b c => (b -> c) -> (a -> b) -> (a -> c)
(#.) _f = coerce
Expand Down
48 changes: 45 additions & 3 deletions libraries/base/Data/Traversable.hs
Expand Up @@ -28,8 +28,10 @@ module Data.Traversable (
-- * Utility functions
for,
forM,
forAccumM,
mapAccumL,
mapAccumR,
mapAccumM,
-- * General definitions for superclass methods
fmapDefault,
foldMapDefault,
Expand Down Expand Up @@ -99,7 +101,7 @@ import Data.Either ( Either(..) )
import Data.Foldable
import Data.Functor
import Data.Functor.Identity ( Identity(..) )
import Data.Functor.Utils ( StateL(..), StateR(..) )
import Data.Functor.Utils ( StateL(..), StateR(..), StateT(..), (#.) )
import Data.Monoid ( Dual(..), Sum(..), Product(..),
First(..), Last(..), Alt(..), Ap(..) )
import Data.Ord ( Down(..) )
Expand Down Expand Up @@ -482,6 +484,45 @@ mapAccumR :: forall t s a b. Traversable t
-- See Note [Function coercion] in Data.Functor.Utils.
mapAccumR f s t = coerce (traverse @t @(StateR s) @a @b) (flip f) t s

-- | The `mapAccumM` function behaves like a combination of `mapM` and
-- `mapAccumL` that traverses the structure while evaluating the actions
-- and passing an accumulating parameter from left to right.
-- It returns a final value of this accumulator together with the new structure.
-- The accummulator is often used for caching the intermediate results of a computation.
--
-- @since 4.18.0.0
--
-- ==== __Examples__
--
-- Basic usage:
--
-- >>> let expensiveDouble a = putStrLn ("Doubling " <> show a) >> pure (2 * a)
-- >>> :{
-- mapAccumM (\cache a -> case lookup a cache of
-- Nothing -> expensiveDouble a >>= \double -> pure ((a, double):cache, double)
-- Just double -> pure (cache, double)
-- ) [] [1, 2, 3, 1, 2, 3]
-- :}
-- Doubling 1
-- Doubling 2
-- Doubling 3
-- ([(3,6),(2,4),(1,2)],[2,4,6,2,4,6])
--
mapAccumM
:: forall m t s a b. (Monad m, Traversable t)
=> (s -> a -> m (s, b))
-> s -> t a -> m (s, t b)
mapAccumM f s t = coerce (mapM @t @(StateT s m) @a @b) (StateT #. flip f) t s

-- | 'forAccumM' is 'mapAccumM' with the arguments rearranged.
--
-- @since 4.18.0.0
forAccumM
:: (Monad m, Traversable t)
=> s -> t a -> (s -> a -> m (s, b)) -> m (s, t b)
{-# INLINE forAccumM #-}
forAccumM s t f = mapAccumM f s t

-- | This function may be used as a value for `fmap` in a `Functor`
-- instance, provided that 'traverse' is defined. (Using
-- `fmapDefault` with a `Traversable` instance defined only by
Expand Down Expand Up @@ -573,8 +614,9 @@ foldMapDefault = coerce (traverse @t @(Const m) @a @())
--
-- When the traversable term is a simple variable or expression, and the
-- monadic action to run is a non-trivial do block, it can be more natural to
-- write the action last. This idiom is supported by 'for' and 'forM', which
-- are the flipped versions of 'traverse' and 'mapM', respectively.
-- write the action last. This idiom is supported by 'for', 'forM', and
-- 'forAccumM' which are the flipped versions of 'traverse', 'mapM', and
-- 'mapAccumM' respectively.

------------------

Expand Down
2 changes: 2 additions & 0 deletions libraries/base/changelog.md
Expand Up @@ -9,6 +9,8 @@
* `Numeric.Natural` re-exports `GHC.Natural.minusNaturalMaybe`.
* Add `Data.Foldable1` and `Data.Bifoldable1`.
* Add `applyWhen` to `Data.Function`.
* Add functions `mapAccumM` and `forAccumM` to `Data.Traversable`, per the
[Core Libraries proposal](https://github.com/haskell/core-libraries-committee/issues/65).

## 4.17.0.0 *TBA*

Expand Down

0 comments on commit 4b98c5c

Please sign in to comment.