Proposal: add traverseWithKey to Data.Map

[Max Bolingbroke]

Hi Haskellers,

I propose to add a new function, traverseWithKey, to Data.Map:

"""
-- | /O(n)/.
-- @'traverseWithKey' f s == 'fmap' 'fromList' ('traverse' (\(k, v) ->
fmap ((,) k) (f k v)) ('toList' s))@
-- That is, behaves exactly like a regular 'traverse' except that the traversing
-- function also has access to the key associated with a value.
--
-- > traverseWithKey (\k v -> if k + v < 10 then Just (v + 1) else
Nothing) (fromList [(1, 2), (5, 4)]) == Just (fromList [(1, 3), (5,
5)])
-- > traverseWithKey (\k v -> if k + v < 10 then Just (v + 1) else
Nothing) (fromList [(5, 5)]) == Nothing
traverseWithKey :: Applicative t => (k -> a -> t b) -> Map k a -> t (Map k b)
"""

This is a rather useful function, and if we define it in SATed style
and with an INLINE pragma (as in my attached patch), GHC can generate
really good code for it at use sites.

While the utility of a traversal function like this is clear, we can
also use it to define other useful combinators such as
mapWithKeyM/mapMWithKey.

It can also be used to define efficient right/left with-key-folds on
Map that avoid a problem in GHC's strictness analyser. For example,
this function:

"""
c m = M.foldlWithKey (\k v (a, b) -> if k + v > 2 then (a, b) else (b,
a)) (0, 1) m
"""

Generates a loop which allocates a new pair on every iteration. (The
situation is the same with foldrWithKey.). If we have traverseWithKey,
we can work around it by defining an optimised fold for accumulators
that are pairs:

"""
newtype State2L s1 s2 a = State2L { unState2L :: s1 -> s2 -> (s1, s2) }

instance Functor (State2L s1 s2) where
    fmap _ = State2L . unState2L

instance Applicative (State2L s1 s2) where
    pure _ = State2L (,)
    mf <*> mx = State2L $ \s1 s2 -> case unState2L mf s1 s2 of (s1,
s2) -> unState2L mx s1 s2 -- NB: left side first

{-# INLINE foldl2WithKey' #-}
foldl2WithKey' :: ((a1, a2) -> k -> v -> (a1, a2)) -> (a1, a2) ->
M.Map k v -> (a1, a2)
foldl2WithKey' f (a1, a2) kvs = unState2L (traverseWithKey (\k v ->
State2L $ \a1 a2 -> f (a1, a2) k v) kvs) a1 a2
"""

The loop resulting from foldl2WithKey' does not allocate any pairs
when compiled with current GHCs.

Of course, traverseWithKey is also sufficiently general that it can
implement many functions already exported from Data.Map (with varying
degrees of efficiency), such as find{Max,Min}, mapWithKey,
fold{r,l}WithKey, withAccumWithKey and mapAccumL.

There is precedence for the "traverseWithKey" name:
 1. "unordered-containers" has a function with exactly this name and
compatible type, used for HashMaps [1]
 2. The "keys" package defines a type class TraversableWithKey with a
compatible type [2]

This is an API addition so breakage from the change should be low.
Overall this should be a low risk addition to the interface which adds
a lot of flexibility. What do you think?

Deadline: 2 weeks (i.e. 28th March)

Patch is attached, assuming it doesn't get stripped by the mailing list manager.

Cheers,
Max

[1] http://hackage.haskell.org/packages/archive/unordered-containers/0.1.1.0/doc/html/src/Data-HashMap-Common.html
[2] http://hackage.haskell.org/packages/archive/keys/0.1.0/doc/html/src/Data-Key.html
-------------- next part --------------
A non-text attachment was scrubbed...
Name: traverseWithKey.patch
Type: application/octet-stream
Size: 2502 bytes
Desc: not available
URL: <http://www.haskell.org/pipermail/libraries/attachments/20120314/f5b63328/attachment.obj>