Add traverseWithKey to Map and IntMap API
authorMax Bolingbroke <batterseapower@hotmail.com>
Fri, 30 Mar 2012 13:20:15 +0000 (14:20 +0100)
committerMax Bolingbroke <batterseapower@hotmail.com>
Fri, 30 Mar 2012 13:56:32 +0000 (14:56 +0100)
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
Data/IntMap/Lazy.hs
Data/IntMap/Strict.hs
Data/Map/Base.hs
Data/Map/Lazy.hs
Data/Map/Strict.hs

index 3d5f923..cdb85a7 100644 (file)
@@ -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.
 --
index e6b2d99..b66ee13 100644 (file)
@@ -113,6 +113,7 @@ module Data.IntMap.Lazy (
             -- ** Map
             , IM.map
             , mapWithKey
+            , traverseWithKey
             , mapAccum
             , mapAccumWithKey
             , mapAccumRWithKey
index 17d6b5d..91508cf 100644 (file)
@@ -117,6 +117,7 @@ module Data.IntMap.Strict (
             -- ** Map
             , map
             , mapWithKey
+            , traverseWithKey
             , mapAccum
             , mapAccumWithKey
             , mapAccumRWithKey
index d9cf4ce..48d3723 100644 (file)
@@ -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
index acca326..020659c 100644 (file)
@@ -109,6 +109,7 @@ module Data.Map.Lazy (
             -- ** Map
             , M.map
             , mapWithKey
+            , traverseWithKey
             , mapAccum
             , mapAccumWithKey
             , mapAccumRWithKey
index e99521d..f8ffc5e 100644 (file)
@@ -129,6 +129,7 @@ module Data.Map.Strict
     -- ** Map
     , map
     , mapWithKey
+    , traverseWithKey
     , mapAccum
     , mapAccumWithKey
     , mapAccumRWithKey