add Data.Foldable.{msum,asum}, plus tweaks to comments

[haskell-directory.git] / Data / Map.hs

diff --git a/Data/Map.hs b/Data/Map.hs
index 269fe5e..da83cf2 100644 (file)
--- a/Data/Map.hs
+++ b/Data/Map.hs
@@ -11,10 +11,11 @@
 --
 -- An efficient implementation of maps from keys to values (dictionaries).
 --
--- This module is intended to be imported @qualified@, to avoid name
--- clashes with Prelude functions.  eg.
+-- Since many function names (but not the type name) clash with
+-- "Prelude" names, this module is usually imported @qualified@, e.g.
 --
--- >  import Data.Map as Map
+-- >  import Data.Map (Map)
+-- >  import qualified Data.Map as Map
 --
 -- The implementation of 'Map' is based on /size balanced/ binary trees (or
 -- trees of /bounded balance/) as described by:
@@ -123,6 +124,11 @@ module Data.Map  (
             , partition
             , partitionWithKey
 
+            , mapMaybe
+            , mapMaybeWithKey
+            , mapEither
+            , mapEitherWithKey
+
             , split         
             , splitLookup   
 
@@ -870,6 +876,33 @@ partitionWithKey p (Bin _ kx x l r)
     (l1,l2) = partitionWithKey p l
     (r1,r2) = partitionWithKey p r
 
+-- | /O(n)/. Map values and collect the 'Just' results.
+mapMaybe :: Ord k => (a -> Maybe b) -> Map k a -> Map k b
+mapMaybe f m
+  = mapMaybeWithKey (\k x -> f x) m
+
+-- | /O(n)/. Map keys\/values and collect the 'Just' results.
+mapMaybeWithKey :: Ord k => (k -> a -> Maybe b) -> Map k a -> Map k b
+mapMaybeWithKey f Tip = Tip
+mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of
+  Just y  -> join kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r)
+  Nothing -> merge (mapMaybeWithKey f l) (mapMaybeWithKey f r)
+
+-- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
+mapEither :: Ord k => (a -> Either b c) -> Map k a -> (Map k b, Map k c)
+mapEither f m
+  = mapEitherWithKey (\k x -> f x) m
+
+-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
+mapEitherWithKey :: Ord k =>
+  (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c)
+mapEitherWithKey f Tip = (Tip, Tip)
+mapEitherWithKey f (Bin _ kx x l r) = case f kx x of
+  Left y  -> (join kx y l1 r1, merge l2 r2)
+  Right z -> (merge l1 r1, join kx z l2 r2)
+  where
+    (l1,l2) = mapEitherWithKey f l
+    (r1,r2) = mapEitherWithKey f r
 
 {--------------------------------------------------------------------
   Mapping