Some small optimisations, generalise the type of unfold
[haskell-directory.git] / Data / IntMap.hs
index f9ee03a..652a23e 100644 (file)
@@ -1,4 +1,4 @@
-{-# OPTIONS -cpp -fglasgow-exts #-} 
+{-# OPTIONS -cpp -fglasgow-exts -fno-bang-patterns #-} 
 -----------------------------------------------------------------------------
 -- Module      :  Data.IntMap
 -- Copyright   :  (c) Daan Leijen 2002
@@ -45,6 +45,7 @@ module Data.IntMap  (
             , null
             , size
             , member
+            , notMember
            , lookup
             , findWithDefault
             
@@ -63,6 +64,7 @@ module Data.IntMap  (
             , update
             , updateWithKey
             , updateLookupWithKey
+            , alter
   
             -- * Combine
 
@@ -136,7 +138,9 @@ import Prelude hiding (lookup,map,filter,foldr,foldl,null)
 import Data.Bits 
 import Data.Int
 import qualified Data.IntSet as IntSet
+import Data.Monoid (Monoid(..))
 import Data.Typeable
+import Data.Foldable (Foldable(foldMap))
 
 {-
 -- just for testing
@@ -210,6 +214,16 @@ type Prefix = Int
 type Mask   = Int
 type Key    = Int
 
+instance Monoid (IntMap a) where
+    mempty  = empty
+    mappend = union
+    mconcat = unions
+
+instance Foldable IntMap where
+    foldMap f Nil = mempty
+    foldMap f (Tip _k v) = f v
+    foldMap f (Bin _ _ l r) = foldMap f l `mappend` foldMap f r
+
 #if __GLASGOW_HASKELL__
 
 {--------------------------------------------------------------------
@@ -224,6 +238,7 @@ instance Data a => Data (IntMap a) where
   toConstr _    = error "toConstr"
   gunfold _ _   = error "gunfold"
   dataTypeOf _  = mkNorepType "Data.IntMap.IntMap"
+  dataCast1 f   = gcast1 f
 
 #endif
 
@@ -250,9 +265,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
@@ -316,11 +340,19 @@ insert k x t
 
 -- right-biased insertion, used by 'union'
 -- | /O(min(n,W))/. Insert with a combining function.
+-- @'insertWith' f key value mp@ 
+-- will insert the pair (key, value) into @mp@ if key does
+-- not exist in the map. If the key does exist, the function will
+-- insert @f new_value old_value@.
 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
 insertWith f k x t
   = insertWithKey (\k x y -> f x y) k x t
 
 -- | /O(min(n,W))/. Insert with a combining function.
+-- @'insertWithKey' f key value mp@ 
+-- will insert the pair (key, value) into @mp@ if key does
+-- not exist in the map. If the key does exist, the function will
+-- insert @f key new_value old_value@.
 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
 insertWithKey f k x t
   = case t of
@@ -420,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
 --------------------------------------------------------------------}
@@ -798,6 +854,20 @@ partitionWithKey pred t
 split :: Key -> IntMap a -> (IntMap a,IntMap a)
 split k t
   = case t of
+      Bin p m l r 
+          | m < 0 -> (if k >= 0 -- handle negative numbers.
+                      then let (lt,gt) = split' k l in (union r lt, gt)
+                      else let (lt,gt) = split' k r in (lt, union gt l))
+          | otherwise   -> split' k t
+      Tip ky y 
+        | k>ky      -> (t,Nil)
+        | k<ky      -> (Nil,t)
+        | otherwise -> (Nil,Nil)
+      Nil -> (Nil,Nil)
+
+split' :: Key -> IntMap a -> (IntMap a,IntMap a)
+split' k t
+  = case t of
       Bin p m l r
         | nomatch k p m -> if k>p then (t,Nil) else (Nil,t)
         | zero k m  -> let (lt,gt) = split k l in (lt,union gt r)
@@ -814,6 +884,20 @@ splitLookup :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
 splitLookup k t
   = case t of
       Bin p m l r
+          | m < 0 -> (if k >= 0 -- handle negative numbers.
+                      then let (lt,found,gt) = splitLookup' k l in (union r lt,found, gt)
+                      else let (lt,found,gt) = splitLookup' k r in (lt,found, union gt l))
+          | otherwise   -> splitLookup' k t
+      Tip ky y 
+        | k>ky      -> (t,Nothing,Nil)
+        | k<ky      -> (Nil,Nothing,t)
+        | otherwise -> (Nil,Just y,Nil)
+      Nil -> (Nil,Nothing,Nil)
+
+splitLookup' :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
+splitLookup' k t
+  = case t of
+      Bin p m l r
         | nomatch k p m -> if k>p then (t,Nothing,Nil) else (Nil,Nothing,t)
         | zero k m  -> let (lt,found,gt) = splitLookup k l in (lt,found,union gt r)
         | otherwise -> let (lt,found,gt) = splitLookup k r in (union l lt,found,gt)
@@ -849,10 +933,20 @@ foldWithKey f z t
 foldr :: (Key -> a -> b -> b) -> b -> IntMap a -> b
 foldr f z t
   = case t of
-      Bin p m l r -> foldr f (foldr f z r) l
+      Bin 0 m l r | m < 0 -> foldr' f (foldr' f z l) r  -- put negative numbers before.
+      Bin _ _ _ _ -> foldr' f z t
+      Tip k x     -> f k x z
+      Nil         -> z
+
+foldr' :: (Key -> a -> b -> b) -> b -> IntMap a -> b
+foldr' f z t
+  = case t of
+      Bin p m l r -> foldr' f (foldr' f z r) l
       Tip k x     -> f k x z
       Nil         -> z
 
+
+
 {--------------------------------------------------------------------
   List variations 
 --------------------------------------------------------------------}