Added 'alter'

Added 'alter :: (Maybe a -> Maybe a) -> k -> Map k a -> Map k a' to IntMap and Map
This addresses ticket #665

2 files changed, 43 insertions, 0 deletions

diff --git a/libraries/base/Data/IntMap.hs b/libraries/base/Data/IntMap.hs
index ca9e4e0934..652a23e565 100644
--- a/libraries/base/Data/IntMap.hs
+++ b/libraries/base/Data/IntMap.hs
@@ -64,6 +64,7 @@ module Data.IntMap  (
             , update
             , updateWithKey
             , updateLookupWithKey
+            , alter
   
             -- * Combine
 
@@ -451,6 +452,30 @@ updateLookupWithKey f k t
       Nil -> (Nothing,Nil)
 
 
+
+-- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
+-- 'alter' can be used to insert, delete, or update a value in a 'Map'.
+-- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@
+alter f k t
+  = case t of
+      Bin p m l r 
+        | nomatch k p m -> case f Nothing of 
+                             Nothing -> t
+                             Just x -> join k (Tip k x) p t
+        | zero k m      -> bin p m (alter f k l) r
+        | otherwise     -> bin p m l (alter f k r)
+      Tip ky y          
+        | k==ky         -> case f (Just y) of
+                             Just x -> Tip ky x
+                             Nothing -> Nil
+        | otherwise     -> case f Nothing of
+                             Just x -> join k (Tip k x) ky t
+                             Nothing -> Tip ky y
+      Nil               -> case f Nothing of
+                             Just x -> Tip k x
+                             Nothing -> Nil
+
+
 {--------------------------------------------------------------------
   Union
 --------------------------------------------------------------------}
diff --git a/libraries/base/Data/Map.hs b/libraries/base/Data/Map.hs
index 92b21a1684..517a6bbbaf 100644
--- a/libraries/base/Data/Map.hs
+++ b/libraries/base/Data/Map.hs
@@ -63,6 +63,7 @@ module Data.Map  (
             , update
             , updateWithKey
             , updateLookupWithKey
+            , alter
 
             -- * Combine
 
@@ -427,6 +428,23 @@ updateLookupWithKey f k t
                        Just x' -> (Just x',Bin sx kx x' l r)
                        Nothing -> (Just x,glue l r)
 
+-- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
+-- 'alter' can be used to insert, delete, or update a value in a 'Map'.
+-- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@
+alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
+alter f k t
+  = case t of
+      Tip -> case f Nothing of
+               Nothing -> Tip
+               Just x -> singleton k x
+      Bin sx kx x l r 
+          -> case compare k kx of
+               LT -> balance kx x (alter f k l) r
+               GT -> balance kx x l (alter f k r)
+               EQ -> case f (Just x) of
+                       Just x' -> Bin sx kx x' l r
+                       Nothing -> glue l r
+
 {--------------------------------------------------------------------
   Indexing
 --------------------------------------------------------------------}