add System.Posix.Types to default nhc98 build
[haskell-directory.git] / Data / IntMap.hs
index 25b1558..9b897a3 100644 (file)
@@ -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
 --
 -- 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   
 
@@ -126,6 +135,23 @@ module Data.IntMap  (
             , isSubmapOf, isSubmapOfBy
             , isProperSubmapOf, isProperSubmapOfBy
             
+            -- * Min\/Max
+
+            , maxView
+            , minView
+            , findMin   
+            , findMax
+            , deleteMin
+            , deleteMax
+            , deleteFindMin
+            , deleteFindMax
+            , updateMin
+            , updateMax
+            , updateMinWithKey
+            , updateMaxWithKey 
+            , minViewWithKey
+            , maxViewWithKey
+
             -- * Debugging
             , showTree
             , showTreeWith
@@ -134,11 +160,12 @@ module Data.IntMap  (
 
 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
 import Data.Bits 
-import Data.Int
-import Data.Monoid
 import qualified Data.IntSet as IntSet
+import Data.Monoid (Monoid(..))
 import Data.Typeable
-
+import Data.Foldable (Foldable(foldMap))
+import Control.Monad ( liftM )
+import Control.Arrow (ArrowZero)
 {-
 -- just for testing
 import qualified Prelude
@@ -148,12 +175,12 @@ import qualified List
 -}  
 
 #if __GLASGOW_HASKELL__
-import Data.Generics.Basics
-import Data.Generics.Instances
+import Text.Read
+import Data.Generics.Basics (Data(..), mkNorepType)
+import Data.Generics.Instances ()
 #endif
 
 #if __GLASGOW_HASKELL__ >= 503
-import GHC.Word
 import GHC.Exts ( Word(..), Int(..), shiftRL# )
 #elif __GLASGOW_HASKELL__
 import Word
@@ -210,6 +237,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 +261,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 +288,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
@@ -297,11 +344,11 @@ singleton k x
 
 {--------------------------------------------------------------------
   Insert
-  'insert' is the inlined version of 'insertWith (\k x y -> x)'
 --------------------------------------------------------------------}
--- | /O(min(n,W))/. Insert a new key\/value pair in the map. When the key 
--- is already an element of the set, its value is replaced by the new value, 
--- ie. 'insert' is left-biased.
+-- | /O(min(n,W))/. Insert a new key\/value pair in the map.
+-- If the key is already present in the map, the associated value is
+-- replaced with the supplied value, i.e. 'insert' is equivalent to
+-- @'insertWith' 'const'@.
 insert :: Key -> a -> IntMap a -> IntMap a
 insert k x t
   = case t of
@@ -316,11 +363,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 +475,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
 --------------------------------------------------------------------}
@@ -433,7 +512,9 @@ unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a
 unionsWith f ts
   = foldlStrict (unionWith f) empty ts
 
--- | /O(n+m)/. The (left-biased) union of two sets. 
+-- | /O(n+m)/. The (left-biased) union of two maps. 
+-- It prefers the first map when duplicate keys are encountered,
+-- i.e. (@'union' == 'unionWith' 'const'@).
 union :: IntMap a -> IntMap a -> IntMap a
 union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
   | shorter m1 m2  = union1
@@ -606,6 +687,115 @@ intersectionWithKey f t Nil = Nil
 
 
 {--------------------------------------------------------------------
+  Min\/Max
+--------------------------------------------------------------------}
+
+-- | /O(log n)/. Update the value at the minimal key.
+updateMinWithKey :: (Key -> a -> a) -> IntMap a -> IntMap a
+updateMinWithKey f t
+    = case t of
+        Bin p m l r | m < 0 -> let t' = updateMinWithKeyUnsigned f l in Bin p m t' r
+        Bin p m l r         -> let t' = updateMinWithKeyUnsigned f r in Bin p m l t'
+        Tip k y -> Tip k (f k y)
+        Nil -> error "maxView: empty map has no maximal element"
+
+updateMinWithKeyUnsigned f t
+    = case t of
+        Bin p m l r -> let t' = updateMinWithKeyUnsigned f r in Bin p m l t'
+        Tip k y -> Tip k (f k y)
+
+-- | /O(log n)/. Update the value at the maximal key.
+updateMaxWithKey :: (Key -> a -> a) -> IntMap a -> IntMap a
+updateMaxWithKey f t
+    = case t of
+        Bin p m l r | m < 0 -> let t' = updateMaxWithKeyUnsigned f r in Bin p m r t'
+        Bin p m l r         -> let t' = updateMaxWithKeyUnsigned f l in Bin p m t' l
+        Tip k y -> Tip k (f k y)
+        Nil -> error "maxView: empty map has no maximal element"
+
+updateMaxWithKeyUnsigned f t
+    = case t of
+        Bin p m l r -> let t' = updateMaxWithKeyUnsigned f r in Bin p m l t'
+        Tip k y -> Tip k (f k y)
+
+
+-- | /O(log n)/. Retrieves the maximal (key,value) couple of the map, and the map stripped from that element.
+-- @fail@s (in the monad) when passed an empty map.
+maxViewWithKey :: (Monad m) => IntMap a -> m ((Key, a), IntMap a)
+maxViewWithKey t
+    = case t of
+        Bin p m l r | m < 0 -> let (result, t') = maxViewUnsigned l in return (result, bin p m t' r)
+        Bin p m l r         -> let (result, t') = maxViewUnsigned r in return (result, bin p m l t')
+        Tip k y -> return ((k,y), Nil)
+        Nil -> fail "maxView: empty map has no maximal element"
+
+maxViewUnsigned t 
+    = case t of
+        Bin p m l r -> let (result,t') = maxViewUnsigned r in (result,bin p m l t')
+        Tip k y -> ((k,y), Nil)
+
+-- | /O(log n)/. Retrieves the minimal (key,value) couple of the map, and the map stripped from that element.
+-- @fail@s (in the monad) when passed an empty map.
+minViewWithKey :: (Monad m) => IntMap a -> m ((Key, a), IntMap a)
+minViewWithKey t
+    = case t of
+        Bin p m l r | m < 0 -> let (result, t') = minViewUnsigned r in return (result, bin p m l t')
+        Bin p m l r         -> let (result, t') = minViewUnsigned l in return (result, bin p m t' r)
+        Tip k y -> return ((k,y),Nil)
+        Nil -> fail "minView: empty map has no minimal element"
+
+minViewUnsigned t 
+    = case t of
+        Bin p m l r -> let (result,t') = minViewUnsigned l in (result,bin p m t' r)
+        Tip k y -> ((k,y),Nil)
+
+
+-- | /O(log n)/. Update the value at the maximal key.
+updateMax :: (a -> a) -> IntMap a -> IntMap a
+updateMax f = updateMaxWithKey (const f)
+
+-- | /O(log n)/. Update the value at the minimal key.
+updateMin :: (a -> a) -> IntMap a -> IntMap a
+updateMin f = updateMinWithKey (const f)
+
+
+-- Duplicate the Identity monad here because base < mtl.
+newtype Identity a = Identity { runIdentity :: a }
+instance Monad Identity where
+       return a = Identity a
+       m >>= k  = k (runIdentity m)
+-- Similar to the Arrow instance.
+first f (x,y) = (f x,y)
+
+
+-- | /O(log n)/. Retrieves the maximal key of the map, and the map stripped from that element.
+-- @fail@s (in the monad) when passed an empty map.
+maxView t = liftM (first snd) (maxViewWithKey t)
+
+-- | /O(log n)/. Retrieves the minimal key of the map, and the map stripped from that element.
+-- @fail@s (in the monad) when passed an empty map.
+minView t = liftM (first snd) (minViewWithKey t)
+
+-- | /O(log n)/. Delete and find the maximal element.
+deleteFindMax = runIdentity . maxView
+
+-- | /O(log n)/. Delete and find the minimal element.
+deleteFindMin = runIdentity . minView
+
+-- | /O(log n)/. The minimal key of the map.
+findMin = fst . runIdentity . minView
+
+-- | /O(log n)/. The maximal key of the map.
+findMax = fst . runIdentity . maxView
+
+-- | /O(log n)/. Delete the minimal key.
+deleteMin = snd . runIdentity . minView
+
+-- | /O(log n)/. Delete the maximal key.
+deleteMax = snd . runIdentity . maxView
+
+
+{--------------------------------------------------------------------
   Submap
 --------------------------------------------------------------------}
 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). 
@@ -789,6 +979,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
@@ -796,7 +1016,22 @@ 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)
         | otherwise -> let (lt,gt) = split k r in (union l lt,gt)
       Tip ky y 
@@ -811,6 +1046,21 @@ 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)
       Tip ky y 
@@ -845,10 +1095,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 
 --------------------------------------------------------------------}
@@ -971,21 +1231,12 @@ instance Functor IntMap where
     fmap = map
 
 {--------------------------------------------------------------------
-  Monoid 
---------------------------------------------------------------------}
-
-instance Ord a => Monoid (IntMap a) where
-    mempty = empty
-    mappend = union
-    mconcat = unions
-
-{--------------------------------------------------------------------
   Show 
 --------------------------------------------------------------------}
 
 instance Show a => Show (IntMap a) where
-  showsPrec d t   = showMap (toList t)
-
+  showsPrec d m   = showParen (d > 10) $
+    showString "fromList " . shows (toList m)
 
 showMap :: (Show a) => [(Key,a)] -> ShowS
 showMap []     
@@ -997,7 +1248,25 @@ showMap (x:xs)
     showTail (x:xs) = showChar ',' . showElem x . showTail xs
     
     showElem (k,x)  = shows k . showString ":=" . shows x
-  
+
+{--------------------------------------------------------------------
+  Read
+--------------------------------------------------------------------}
+instance (Read e) => Read (IntMap e) where
+#ifdef __GLASGOW_HASKELL__
+  readPrec = parens $ prec 10 $ do
+    Ident "fromList" <- lexP
+    xs <- readPrec
+    return (fromList xs)
+
+  readListPrec = readListPrecDefault
+#else
+  readsPrec p = readParen (p > 10) $ \ r -> do
+    ("fromList",s) <- lex r
+    (xs,t) <- reads s
+    return (fromList xs,t)
+#endif
+
 {--------------------------------------------------------------------
   Typeable
 --------------------------------------------------------------------}