[commit: containers] master: Add traverseWithKey to Map and IntMap API (76a2b9d)

Repository : ssh://darcs.haskell.org//srv/darcs/packages/containers

On branch  : master

http://hackage.haskell.org/trac/ghc/changeset/76a2b9d97429497fe4f88cc47eb4945ba7eda163

>---------------------------------------------------------------

commit 76a2b9d97429497fe4f88cc47eb4945ba7eda163
Author: Max Bolingbroke <[email protected]>
Date:   Fri Mar 30 14:20:15 2012 +0100

    Add traverseWithKey to Map and IntMap API
    
    Proposal reviewed and approved by the libraries list.
    Particular thanks goes to Thomas Schilling for his suggestions
    regarding how the function should be documented.

>---------------------------------------------------------------

 Data/IntMap/Base.hs   |   20 +++++++++++++++++---
 Data/IntMap/Lazy.hs   |    1 +
 Data/IntMap/Strict.hs |    1 +
 Data/Map/Base.hs      |   20 +++++++++++++++++---
 Data/Map/Lazy.hs      |    1 +
 Data/Map/Strict.hs    |    1 +
 6 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/Data/IntMap/Base.hs b/Data/IntMap/Base.hs
index 3d5f923..cdb85a7 100644
--- a/Data/IntMap/Base.hs
+++ b/Data/IntMap/Base.hs
@@ -81,6 +81,7 @@ module Data.IntMap.Base (
             -- ** Map
             , map
             , mapWithKey
+            , traverseWithKey
             , mapAccum
             , mapAccumWithKey
             , mapAccumRWithKey
@@ -289,9 +290,7 @@ instance Foldable.Foldable IntMap where
   foldMap f (Bin _ _ l r) = Foldable.foldMap f l `mappend` Foldable.foldMap f r
 
 instance Traversable IntMap where
-    traverse _ Nil = pure Nil
-    traverse f (Tip k v) = Tip k <$> f v
-    traverse f (Bin p m l r) = Bin p m <$> traverse f l <*> traverse f r
+    traverse f = traverseWithKey (\_ -> f)
 
 instance NFData a => NFData (IntMap a) where
     rnf Nil = ()
@@ -1121,6 +1120,21 @@ mapWithKey f t
       Tip k x     -> Tip k (f k x)
       Nil         -> Nil
 
+-- | /O(n)/.
+-- @'traverseWithKey' f s == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@
+-- 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 odd k then Just (succ v) else Nothing) (fromList [(1, 'a'), (5, 'e')]) == Just (fromList [(1, 'b'), (5, 'f')])
+-- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(2, 'c')])           == Nothing
+{-# INLINE traverseWithKey #-}
+traverseWithKey :: Applicative t => (Key -> a -> t b) -> IntMap a -> t (IntMap b)
+traverseWithKey f = go
+  where
+    go Nil = pure Nil
+    go (Tip k v) = Tip k <$> f k v
+    go (Bin p m l r) = Bin p m <$> go l <*> go r
+
 -- | /O(n)/. The function @'mapAccum'@ threads an accumulating
 -- argument through the map in ascending order of keys.
 --
diff --git a/Data/IntMap/Lazy.hs b/Data/IntMap/Lazy.hs
index e6b2d99..b66ee13 100644
--- a/Data/IntMap/Lazy.hs
+++ b/Data/IntMap/Lazy.hs
@@ -113,6 +113,7 @@ module Data.IntMap.Lazy (
             -- ** Map
             , IM.map
             , mapWithKey
+            , traverseWithKey
             , mapAccum
             , mapAccumWithKey
             , mapAccumRWithKey
diff --git a/Data/IntMap/Strict.hs b/Data/IntMap/Strict.hs
index 17d6b5d..91508cf 100644
--- a/Data/IntMap/Strict.hs
+++ b/Data/IntMap/Strict.hs
@@ -117,6 +117,7 @@ module Data.IntMap.Strict (
             -- ** Map
             , map
             , mapWithKey
+            , traverseWithKey
             , mapAccum
             , mapAccumWithKey
             , mapAccumRWithKey
diff --git a/Data/Map/Base.hs b/Data/Map/Base.hs
index d9cf4ce..48d3723 100644
--- a/Data/Map/Base.hs
+++ b/Data/Map/Base.hs
@@ -114,6 +114,7 @@ module Data.Map.Base (
             -- ** Map
             , map
             , mapWithKey
+            , traverseWithKey
             , mapAccum
             , mapAccumWithKey
             , mapAccumRWithKey
@@ -1445,6 +1446,21 @@ mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
 mapWithKey _ Tip = Tip
 mapWithKey f (Bin sx kx x l r) = Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey f r)
 
+-- | /O(n)/.
+-- @'traverseWithKey' f s == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@
+-- 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 odd k then Just (succ v) else Nothing) (fromList [(1, 'a'), (5, 'e')]) == Just (fromList [(1, 'b'), (5, 'f')])
+-- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(2, 'c')])           == Nothing
+{-# INLINE traverseWithKey #-}
+traverseWithKey :: Applicative t => (k -> a -> t b) -> Map k a -> t (Map k b)
+traverseWithKey f = go
+  where
+    go Tip = pure Tip
+    go (Bin s k v l r)
+      = flip (Bin s k) <$> go l <*> f k v <*> go r
+
 -- | /O(n)/. The function 'mapAccum' threads an accumulating
 -- argument through the map in ascending order of keys.
 --
@@ -2326,9 +2342,7 @@ instance Functor (Map k) where
   fmap f m  = map f m
 
 instance Traversable (Map k) where
-  traverse _ Tip = pure Tip
-  traverse f (Bin s k v l r)
-    = flip (Bin s k) <$> traverse f l <*> f v <*> traverse f r
+  traverse f = traverseWithKey (\_ -> f)
 
 instance Foldable.Foldable (Map k) where
   fold Tip = mempty
diff --git a/Data/Map/Lazy.hs b/Data/Map/Lazy.hs
index acca326..020659c 100644
--- a/Data/Map/Lazy.hs
+++ b/Data/Map/Lazy.hs
@@ -109,6 +109,7 @@ module Data.Map.Lazy (
             -- ** Map
             , M.map
             , mapWithKey
+            , traverseWithKey
             , mapAccum
             , mapAccumWithKey
             , mapAccumRWithKey
diff --git a/Data/Map/Strict.hs b/Data/Map/Strict.hs
index e99521d..f8ffc5e 100644
--- a/Data/Map/Strict.hs
+++ b/Data/Map/Strict.hs
@@ -129,6 +129,7 @@ module Data.Map.Strict
     -- ** Map
     , map
     , mapWithKey
+    , traverseWithKey
     , mapAccum
     , mapAccumWithKey
     , mapAccumRWithKey