X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=Data%2FIntMap.hs;h=e4377e5797f0393c6e4fe74f074635a721902433;hb=7a73aaf70fefb4b30c9159f5d15035f8e9c6e114;hp=16064132ca8f8e44046de745f9b27a5141f7f0bb;hpb=94dc10bc5f7cbf36814e9158979d0e8fde7c5eef;p=haskell-directory.git

diff --git a/Data/IntMap.hs b/Data/IntMap.hs
index 1606413..e4377e5 100644
--- a/Data/IntMap.hs
+++ b/Data/IntMap.hs
@@ -1,5 +1,6 @@
-{-# OPTIONS -cpp -fglasgow-exts #-} 
+{-# OPTIONS -cpp -fglasgow-exts -fno-bang-patterns #-} 
 -----------------------------------------------------------------------------
+-- |
 -- Module      :  Data.IntMap
 -- Copyright   :  (c) Daan Leijen 2002
 -- License     :  BSD-style
@@ -9,10 +10,11 @@
 --
 -- An efficient implementation of maps from integer keys to values.
 --
--- 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.IntMap as Map
+-- >  import Data.IntMap (IntMap)
+-- >  import qualified Data.IntMap as IntMap
 --
 -- The implementation is based on /big-endian patricia trees/.  This data
 -- structure performs especially well on binary operations like 'union'
@@ -45,6 +47,7 @@ module Data.IntMap  (
             , null
             , size
             , member
+            , notMember
 	    , lookup
             , findWithDefault
             
@@ -63,6 +66,7 @@ module Data.IntMap  (
             , update
             , updateWithKey
             , updateLookupWithKey
+            , alter
   
             -- * Combine
 
@@ -119,6 +123,11 @@ module Data.IntMap  (
             , partition
             , partitionWithKey
 
+            , mapMaybe
+            , mapMaybeWithKey
+            , mapEither
+            , mapEitherWithKey
+
             , split         
             , splitLookup   
 
@@ -263,9 +272,18 @@ member k m
       Nothing -> False
       Just x  -> True
     
+-- | /O(log n)/. Is the key not a member of the map?
+notMember :: Key -> IntMap a -> Bool
+notMember k m = not $ member k m
+
 -- | /O(min(n,W))/. Lookup the value at a key in the map.
-lookup :: Key -> IntMap a -> Maybe a
-lookup k t
+lookup :: (Monad m) => Key -> IntMap a -> m a
+lookup k t = case lookup' k t of
+    Just x -> return x
+    Nothing -> fail "Data.IntMap.lookup: Key not found"
+
+lookup' :: Key -> IntMap a -> Maybe a
+lookup' k t
   = let nk = natFromInt k  in seq nk (lookupN nk t)
 
 lookupN :: Nat -> IntMap a -> Maybe a
@@ -441,6 +459,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
 --------------------------------------------------------------------}
@@ -812,6 +854,36 @@ partitionWithKey pred t
         | otherwise -> (Nil,t)
       Nil -> (Nil,Nil)
 
+-- | /O(n)/. Map values and collect the 'Just' results.
+mapMaybe :: (a -> Maybe b) -> IntMap a -> IntMap b
+mapMaybe f m
+  = mapMaybeWithKey (\k x -> f x) m
+
+-- | /O(n)/. Map keys\/values and collect the 'Just' results.
+mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b
+mapMaybeWithKey f (Bin p m l r)
+  = bin p m (mapMaybeWithKey f l) (mapMaybeWithKey f r)
+mapMaybeWithKey f (Tip k x) = case f k x of
+  Just y  -> Tip k y
+  Nothing -> Nil
+mapMaybeWithKey f Nil = Nil
+
+-- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
+mapEither :: (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
+mapEither f m
+  = mapEitherWithKey (\k x -> f x) m
+
+-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
+mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
+mapEitherWithKey f (Bin p m l r)
+  = (bin p m l1 r1, bin p m l2 r2)
+  where
+    (l1,l2) = mapEitherWithKey f l
+    (r1,r2) = mapEitherWithKey f r
+mapEitherWithKey f (Tip k x) = case f k x of
+  Left y  -> (Tip k y, Nil)
+  Right z -> (Nil, Tip k z)
+mapEitherWithKey f Nil = (Nil, Nil)
 
 -- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@
 -- where all keys in @map1@ are lower than @k@ and all keys in