Skip to content

Commit

Permalink
Add foldlWithKeyM and foldrWithKeyM
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewthad committed Aug 17, 2023
1 parent e74ac0b commit f21b72c
Showing 1 changed file with 32 additions and 0 deletions.
32 changes: 32 additions & 0 deletions containers/src/Data/Map/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -259,6 +259,10 @@ module Data.Map.Internal (
, foldrWithKey'
, foldlWithKey'

-- ** Monadic folds
, foldlWithKeyM
, foldrWithKeyM

-- * Conversion
, elems
, keys
Expand Down Expand Up @@ -3315,6 +3319,34 @@ foldlWithKey' f z = go z
in go (f z'' kx x) r
{-# INLINE foldlWithKey' #-}

-- | /O(n)/. Monadic variant of 'foldlWithKey'.
--
-- > 'foldlWithKeyM\'' f z0 (fromList [(0,'a'),(1,'b'),(2,'c')]) =
-- > f z0 0 'a' >>= (\z1 -> f z1 1 'b' >>= (\z2 -> f z2 2 'c'))
foldlWithKeyM :: Monad m => (a -> k -> b -> m a) -> a -> Map k b -> m a
foldlWithKeyM f z = go z
where
go z' Tip = return z'
go z' (Bin _ kx x l r) = do
z'' <- go z' l
z''' <- f z'' kx x
go z''' r
{-# INLINE foldlWithKeyM #-}

-- | /O(n)/. Monadic variant of 'foldrWithKey'.
--
-- > 'foldrWithKeyM\'' f z0 (fromList [(0,'a'),(1,'b'),(2,'c')]) =
-- > (f 2 'c' >=> f 1 'b' >=> f 0 'a') z0
foldrWithKeyM :: Monad m => (k -> a -> b -> m b) -> b -> Map k a -> m b
foldrWithKeyM f z = go z
where
go z' Tip = return z'
go z' (Bin _ kx x l r) = do
z'' <- go z' r
z''' <- f kx x z''
go z''' l
{-# INLINE foldrWithKeyM #-}

-- | \(O(n)\). Fold the keys and values in the map using the given monoid, such that
--
-- @'foldMapWithKey' f = 'Prelude.fold' . 'mapWithKey' f@
Expand Down

0 comments on commit f21b72c

Please sign in to comment.