[project @ 2005-11-29 14:31:59 by ross]
[haskell-directory.git] / Data / IntMap.hs
index 46c6bee..1be0bbe 100644 (file)
@@ -135,9 +135,10 @@ 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))
 
 {-
 -- just for testing
@@ -148,6 +149,7 @@ import qualified List
 -}  
 
 #if __GLASGOW_HASKELL__
+import Text.Read
 import Data.Generics.Basics
 import Data.Generics.Instances
 #endif
@@ -188,7 +190,8 @@ shiftRL x i   = shiftR x i
   Operators
 --------------------------------------------------------------------}
 
--- | /O(min(n,W))/. Find the value of a key. Calls @error@ when the element can not be found.
+-- | /O(min(n,W))/. Find the value at a key.
+-- Calls 'error' when the element can not be found.
 
 (!) :: IntMap a -> Key -> a
 m ! k    = find' k m
@@ -209,12 +212,22 @@ type Prefix = Int
 type Mask   = Int
 type Key    = Int
 
+instance Ord a => 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__
+
 {--------------------------------------------------------------------
   A Data instance  
 --------------------------------------------------------------------}
 
-#if __GLASGOW_HASKELL__
-
 -- This instance preserves data abstraction at the cost of inefficiency.
 -- We omit reflection services for the sake of data abstraction.
 
@@ -249,7 +262,7 @@ member k m
       Nothing -> False
       Just x  -> True
     
--- | /O(min(n,W))/. Lookup the value of a key in the map.
+-- | /O(min(n,W))/. Lookup the value at a key in the map.
 lookup :: Key -> IntMap a -> Maybe a
 lookup k t
   = let nk = natFromInt k  in seq nk (lookupN nk t)
@@ -272,8 +285,9 @@ find' k m
       Just x  -> x
 
 
--- | /O(min(n,W))/. The expression @(findWithDefault def k map)@ returns the value of key @k@ or returns @def@ when
--- the key is not an element of the map.
+-- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
+-- returns the value at key @k@ or returns @def@ when the key is not an
+-- element of the map.
 findWithDefault :: a -> Key -> IntMap a -> a
 findWithDefault def k m
   = case lookup k m of
@@ -295,11 +309,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
@@ -314,11 +328,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
@@ -332,9 +354,9 @@ insertWithKey f k x t
       Nil -> Tip k x
 
 
--- | /O(min(n,W))/. The expression (@insertLookupWithKey f k x map@) is a pair where
--- the first element is equal to (@lookup k map@) and the second element
--- equal to (@insertWithKey f k x map@).
+-- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
+-- is a pair where the first element is equal to (@'lookup' k map@)
+-- and the second element equal to (@'insertWithKey' f k x map@).
 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
 insertLookupWithKey f k x t
   = case t of
@@ -378,16 +400,16 @@ adjustWithKey ::  (Key -> a -> a) -> Key -> IntMap a -> IntMap a
 adjustWithKey f k m
   = updateWithKey (\k x -> Just (f k x)) k m
 
--- | /O(min(n,W))/. The expression (@update f k map@) updates the value @x@
--- at @k@ (if it is in the map). If (@f x@) is @Nothing@, the element is
--- deleted. If it is (@Just y@), the key @k@ is bound to the new value @y@.
+-- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
+-- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
+-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
 update ::  (a -> Maybe a) -> Key -> IntMap a -> IntMap a
 update f k m
   = updateWithKey (\k x -> f x) k m
 
--- | /O(min(n,W))/. The expression (@update f k map@) updates the value @x@
--- at @k@ (if it is in the map). If (@f k x@) is @Nothing@, the element is
--- deleted. If it is (@Just y@), the key @k@ is bound to the new value @y@.
+-- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
+-- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
+-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
 updateWithKey ::  (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
 updateWithKey f k t
   = case t of
@@ -431,7 +453,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
@@ -512,8 +536,8 @@ differenceWith f m1 m2
 
 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
 -- encountered, the combining function is applied to the key and both values.
--- If it returns @Nothing@, the element is discarded (proper set difference). If
--- it returns (@Just y@), the element is updated with a new value @y@. 
+-- If it returns 'Nothing', the element is discarded (proper set difference).
+-- If it returns (@'Just' y@), the element is updated with a new value @y@. 
 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
 differenceWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
   | shorter m1 m2  = difference1
@@ -607,22 +631,22 @@ intersectionWithKey f t Nil = Nil
   Submap
 --------------------------------------------------------------------}
 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). 
--- Defined as (@isProperSubmapOf = isProperSubmapOfBy (==)@).
+-- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
 isProperSubmapOf m1 m2
   = isProperSubmapOfBy (==) m1 m2
 
 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
- The expression (@isProperSubmapOfBy f m1 m2@) returns @True@ when
+ The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
  @m1@ and @m2@ are not equal,
- all keys in @m1@ are in @m2@, and when @f@ returns @True@ when
+ all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
  applied to their respective values. For example, the following 
- expressions are all @True@.
+ expressions are all 'True':
  
   > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
   > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
 
- But the following are all @False@:
+ But the following are all 'False':
  
   > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
   > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
@@ -660,22 +684,23 @@ submapCmp pred (Tip k x) t
 submapCmp pred Nil Nil = EQ
 submapCmp pred Nil t   = LT
 
--- | /O(n+m)/. Is this a submap? Defined as (@isSubmapOf = isSubmapOfBy (==)@).
+-- | /O(n+m)/. Is this a submap?
+-- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
 isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
 isSubmapOf m1 m2
   = isSubmapOfBy (==) m1 m2
 
 {- | /O(n+m)/. 
- The expression (@isSubmapOfBy f m1 m2@) returns @True@ if
- all keys in @m1@ are in @m2@, and when @f@ returns @True@ when
+ The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if
+ all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
  applied to their respective values. For example, the following 
- expressions are all @True@.
+ expressions are all 'True':
  
   > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
   > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
   > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
 
- But the following are all @False@:
+ But the following are all 'False':
  
   > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)])
   > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
@@ -710,20 +735,20 @@ mapWithKey f t
       Tip k x     -> Tip k (f k x)
       Nil         -> Nil
 
--- | /O(n)/. The function @mapAccum@ threads an accumulating
--- argument through the map in an unspecified order.
+-- | /O(n)/. The function @'mapAccum'@ threads an accumulating
+-- argument through the map in ascending order of keys.
 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
 mapAccum f a m
   = mapAccumWithKey (\a k x -> f a x) a m
 
--- | /O(n)/. The function @mapAccumWithKey@ threads an accumulating
--- argument through the map in an unspecified order.
+-- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating
+-- argument through the map in ascending order of keys.
 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
 mapAccumWithKey f a t
   = mapAccumL f a t
 
--- | /O(n)/. The function @mapAccumL@ threads an accumulating
--- argument through the map in pre-order.
+-- | /O(n)/. The function @'mapAccumL'@ threads an accumulating
+-- argument through the map in ascending order of keys.
 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
 mapAccumL f a t
   = case t of
@@ -734,8 +759,8 @@ mapAccumL f a t
       Nil         -> (a,Nil)
 
 
--- | /O(n)/. The function @mapAccumR@ threads an accumulating
--- argument throught the map in post-order.
+-- | /O(n)/. The function @'mapAccumR'@ threads an accumulating
+-- argument throught the map in descending order of keys.
 mapAccumR :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
 mapAccumR f a t
   = case t of
@@ -787,13 +812,28 @@ partitionWithKey pred t
       Nil -> (Nil,Nil)
 
 
--- | /O(log n)/. The expression (@split k map@) is a pair @(map1,map2)@
+-- | /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
 -- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
 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 
@@ -804,32 +844,52 @@ split k t
 
 -- | /O(log n)/. Performs a 'split' but also returns whether the pivot
 -- key was found in the original map.
-splitLookup :: Key -> IntMap a -> (Maybe a,IntMap a,IntMap a)
+splitLookup :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
 splitLookup k t
   = case t of
       Bin p m l r
-        | zero k m  -> let (found,lt,gt) = splitLookup k l in (found,lt,union gt r)
-        | otherwise -> let (found,lt,gt) = splitLookup k r in (found,union l lt,gt)
+          | 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      -> (Nothing,t,Nil)
-        | k<ky      -> (Nothing,Nil,t)
-        | otherwise -> (Just y,Nil,Nil)
-      Nil -> (Nothing,Nil,Nil)
+        | 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 
+        | k>ky      -> (t,Nothing,Nil)
+        | k<ky      -> (Nil,Nothing,t)
+        | otherwise -> (Nil,Just y,Nil)
+      Nil -> (Nil,Nothing,Nil)
 
 {--------------------------------------------------------------------
   Fold
 --------------------------------------------------------------------}
--- | /O(n)/. Fold over the elements of a map in an unspecified order.
+-- | /O(n)/. Fold the values in the map, such that
+-- @'fold' f z == 'Prelude.foldr' f z . 'elems'@.
+-- For example,
 --
--- > sum map   = fold (+) 0 map
 -- > elems map = fold (:) [] map
+--
 fold :: (a -> b -> b) -> b -> IntMap a -> b
 fold f z t
   = foldWithKey (\k x y -> f x y) z t
 
--- | /O(n)/. Fold over the elements of a map in an unspecified order.
+-- | /O(n)/. Fold the keys and values in the map, such that
+-- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
+-- For example,
 --
 -- > keys map = foldWithKey (\k x ks -> k:ks) [] map
+--
 foldWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
 foldWithKey f z t
   = foldr f z t
@@ -837,19 +897,30 @@ 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 
 --------------------------------------------------------------------}
--- | /O(n)/. Return all elements of the map.
+-- | /O(n)/.
+-- Return all elements of the map in the ascending order of their keys.
 elems :: IntMap a -> [a]
 elems m
   = foldWithKey (\k x xs -> x:xs) [] m  
 
--- | /O(n)/. Return all keys of the map.
+-- | /O(n)/. Return all keys of the map in ascending order.
 keys  :: IntMap a -> [Key]
 keys m
   = foldWithKey (\k x ks -> k:ks) [] m
@@ -859,7 +930,7 @@ keysSet :: IntMap a -> IntSet.IntSet
 keysSet m = IntSet.fromDistinctAscList (keys m)
 
 
--- | /O(n)/. Return all key\/value pairs in the map.
+-- | /O(n)/. Return all key\/value pairs in the map in ascending key order.
 assocs :: IntMap a -> [(Key,a)]
 assocs m
   = toList m
@@ -962,21 +1033,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 []     
@@ -988,7 +1050,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
 --------------------------------------------------------------------}
@@ -1006,10 +1086,10 @@ showTree s
   = showTreeWith True False s
 
 
-{- | /O(n)/. The expression (@showTreeWith hang wide map@) shows
+{- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows
  the tree that implements the map. If @hang@ is
- @True@, a /hanging/ tree is shown otherwise a rotated tree is shown. If
- @wide@ is true, an extra wide version is shown.
+ 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
+ @wide@ is 'True', an extra wide version is shown.
 -}
 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
 showTreeWith hang wide t