
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... [2] http://hackage.haskell.org/packages/archive/keys/0.1.0/doc/html/src/Data-Key...