1 {-# OPTIONS -cpp -fglasgow-exts -fno-bang-patterns #-}
2 -----------------------------------------------------------------------------
4 -- Module : Data.IntMap
5 -- Copyright : (c) Daan Leijen 2002
7 -- Maintainer : libraries@haskell.org
8 -- Stability : provisional
9 -- Portability : portable
11 -- An efficient implementation of maps from integer keys to values.
13 -- Since many function names (but not the type name) clash with
14 -- "Prelude" names, this module is usually imported @qualified@, e.g.
16 -- > import Data.IntMap (IntMap)
17 -- > import qualified Data.IntMap as IntMap
19 -- The implementation is based on /big-endian patricia trees/. This data
20 -- structure performs especially well on binary operations like 'union'
21 -- and 'intersection'. However, my benchmarks show that it is also
22 -- (much) faster on insertions and deletions when compared to a generic
23 -- size-balanced map implementation (see "Data.Map" and "Data.FiniteMap").
25 -- * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\",
26 -- Workshop on ML, September 1998, pages 77-86,
27 -- <http://www.cse.ogi.edu/~andy/pub/finite.htm>
29 -- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve
30 -- Information Coded In Alphanumeric/\", Journal of the ACM, 15(4),
31 -- October 1968, pages 514-534.
33 -- Many operations have a worst-case complexity of /O(min(n,W))/.
34 -- This means that the operation can become linear in the number of
35 -- elements with a maximum of /W/ -- the number of bits in an 'Int'
37 -----------------------------------------------------------------------------
41 IntMap, Key -- instance Eq,Show
60 , insertWith, insertWithKey, insertLookupWithKey
118 , fromDistinctAscList
135 , isSubmapOf, isSubmapOfBy
136 , isProperSubmapOf, isProperSubmapOfBy
144 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
146 import qualified Data.IntSet as IntSet
147 import Data.Monoid (Monoid(..))
149 import Data.Foldable (Foldable(foldMap))
153 import qualified Prelude
154 import Debug.QuickCheck
155 import List (nub,sort)
156 import qualified List
159 #if __GLASGOW_HASKELL__
161 import Data.Generics.Basics (Data(..), mkNorepType)
162 import Data.Generics.Instances ()
165 #if __GLASGOW_HASKELL__ >= 503
166 import GHC.Exts ( Word(..), Int(..), shiftRL# )
167 #elif __GLASGOW_HASKELL__
169 import GlaExts ( Word(..), Int(..), shiftRL# )
174 infixl 9 \\{-This comment teaches CPP correct behaviour -}
176 -- A "Nat" is a natural machine word (an unsigned Int)
179 natFromInt :: Key -> Nat
180 natFromInt i = fromIntegral i
182 intFromNat :: Nat -> Key
183 intFromNat w = fromIntegral w
185 shiftRL :: Nat -> Key -> Nat
186 #if __GLASGOW_HASKELL__
187 {--------------------------------------------------------------------
188 GHC: use unboxing to get @shiftRL@ inlined.
189 --------------------------------------------------------------------}
190 shiftRL (W# x) (I# i)
193 shiftRL x i = shiftR x i
196 {--------------------------------------------------------------------
198 --------------------------------------------------------------------}
200 -- | /O(min(n,W))/. Find the value at a key.
201 -- Calls 'error' when the element can not be found.
203 (!) :: IntMap a -> Key -> a
206 -- | /O(n+m)/. See 'difference'.
207 (\\) :: IntMap a -> IntMap b -> IntMap a
208 m1 \\ m2 = difference m1 m2
210 {--------------------------------------------------------------------
212 --------------------------------------------------------------------}
213 -- | A map of integers to values @a@.
215 | Tip {-# UNPACK #-} !Key a
216 | Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(IntMap a) !(IntMap a)
222 instance Monoid (IntMap a) where
227 instance Foldable IntMap where
228 foldMap f Nil = mempty
229 foldMap f (Tip _k v) = f v
230 foldMap f (Bin _ _ l r) = foldMap f l `mappend` foldMap f r
232 #if __GLASGOW_HASKELL__
234 {--------------------------------------------------------------------
236 --------------------------------------------------------------------}
238 -- This instance preserves data abstraction at the cost of inefficiency.
239 -- We omit reflection services for the sake of data abstraction.
241 instance Data a => Data (IntMap a) where
242 gfoldl f z im = z fromList `f` (toList im)
243 toConstr _ = error "toConstr"
244 gunfold _ _ = error "gunfold"
245 dataTypeOf _ = mkNorepType "Data.IntMap.IntMap"
246 dataCast1 f = gcast1 f
250 {--------------------------------------------------------------------
252 --------------------------------------------------------------------}
253 -- | /O(1)/. Is the map empty?
254 null :: IntMap a -> Bool
258 -- | /O(n)/. Number of elements in the map.
259 size :: IntMap a -> Int
262 Bin p m l r -> size l + size r
266 -- | /O(min(n,W))/. Is the key a member of the map?
267 member :: Key -> IntMap a -> Bool
273 -- | /O(log n)/. Is the key not a member of the map?
274 notMember :: Key -> IntMap a -> Bool
275 notMember k m = not $ member k m
277 -- | /O(min(n,W))/. Lookup the value at a key in the map.
278 lookup :: (Monad m) => Key -> IntMap a -> m a
279 lookup k t = case lookup' k t of
281 Nothing -> fail "Data.IntMap.lookup: Key not found"
283 lookup' :: Key -> IntMap a -> Maybe a
285 = let nk = natFromInt k in seq nk (lookupN nk t)
287 lookupN :: Nat -> IntMap a -> Maybe a
291 | zeroN k (natFromInt m) -> lookupN k l
292 | otherwise -> lookupN k r
294 | (k == natFromInt kx) -> Just x
295 | otherwise -> Nothing
298 find' :: Key -> IntMap a -> a
301 Nothing -> error ("IntMap.find: key " ++ show k ++ " is not an element of the map")
305 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
306 -- returns the value at key @k@ or returns @def@ when the key is not an
307 -- element of the map.
308 findWithDefault :: a -> Key -> IntMap a -> a
309 findWithDefault def k m
314 {--------------------------------------------------------------------
316 --------------------------------------------------------------------}
317 -- | /O(1)/. The empty map.
322 -- | /O(1)/. A map of one element.
323 singleton :: Key -> a -> IntMap a
327 {--------------------------------------------------------------------
329 --------------------------------------------------------------------}
330 -- | /O(min(n,W))/. Insert a new key\/value pair in the map.
331 -- If the key is already present in the map, the associated value is
332 -- replaced with the supplied value, i.e. 'insert' is equivalent to
333 -- @'insertWith' 'const'@.
334 insert :: Key -> a -> IntMap a -> IntMap a
338 | nomatch k p m -> join k (Tip k x) p t
339 | zero k m -> Bin p m (insert k x l) r
340 | otherwise -> Bin p m l (insert k x r)
343 | otherwise -> join k (Tip k x) ky t
346 -- right-biased insertion, used by 'union'
347 -- | /O(min(n,W))/. Insert with a combining function.
348 -- @'insertWith' f key value mp@
349 -- will insert the pair (key, value) into @mp@ if key does
350 -- not exist in the map. If the key does exist, the function will
351 -- insert @f new_value old_value@.
352 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
354 = insertWithKey (\k x y -> f x y) k x t
356 -- | /O(min(n,W))/. Insert with a combining function.
357 -- @'insertWithKey' f key value mp@
358 -- will insert the pair (key, value) into @mp@ if key does
359 -- not exist in the map. If the key does exist, the function will
360 -- insert @f key new_value old_value@.
361 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
362 insertWithKey f k x t
365 | nomatch k p m -> join k (Tip k x) p t
366 | zero k m -> Bin p m (insertWithKey f k x l) r
367 | otherwise -> Bin p m l (insertWithKey f k x r)
369 | k==ky -> Tip k (f k x y)
370 | otherwise -> join k (Tip k x) ky t
374 -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
375 -- is a pair where the first element is equal to (@'lookup' k map@)
376 -- and the second element equal to (@'insertWithKey' f k x map@).
377 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
378 insertLookupWithKey f k x t
381 | nomatch k p m -> (Nothing,join k (Tip k x) p t)
382 | zero k m -> let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r)
383 | otherwise -> let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r')
385 | k==ky -> (Just y,Tip k (f k x y))
386 | otherwise -> (Nothing,join k (Tip k x) ky t)
387 Nil -> (Nothing,Tip k x)
390 {--------------------------------------------------------------------
392 [delete] is the inlined version of [deleteWith (\k x -> Nothing)]
393 --------------------------------------------------------------------}
394 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
395 -- a member of the map, the original map is returned.
396 delete :: Key -> IntMap a -> IntMap a
401 | zero k m -> bin p m (delete k l) r
402 | otherwise -> bin p m l (delete k r)
408 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
409 -- a member of the map, the original map is returned.
410 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
412 = adjustWithKey (\k x -> f x) k m
414 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
415 -- a member of the map, the original map is returned.
416 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
418 = updateWithKey (\k x -> Just (f k x)) k m
420 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
421 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
422 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
423 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
425 = updateWithKey (\k x -> f x) k m
427 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
428 -- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
429 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
430 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
435 | zero k m -> bin p m (updateWithKey f k l) r
436 | otherwise -> bin p m l (updateWithKey f k r)
438 | k==ky -> case (f k y) of
444 -- | /O(min(n,W))/. Lookup and update.
445 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
446 updateLookupWithKey f k t
449 | nomatch k p m -> (Nothing,t)
450 | zero k m -> let (found,l') = updateLookupWithKey f k l in (found,bin p m l' r)
451 | otherwise -> let (found,r') = updateLookupWithKey f k r in (found,bin p m l r')
453 | k==ky -> case (f k y) of
454 Just y' -> (Just y,Tip ky y')
455 Nothing -> (Just y,Nil)
456 | otherwise -> (Nothing,t)
461 -- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
462 -- 'alter' can be used to insert, delete, or update a value in a 'Map'.
463 -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@
467 | nomatch k p m -> case f Nothing of
469 Just x -> join k (Tip k x) p t
470 | zero k m -> bin p m (alter f k l) r
471 | otherwise -> bin p m l (alter f k r)
473 | k==ky -> case f (Just y) of
476 | otherwise -> case f Nothing of
477 Just x -> join k (Tip k x) ky t
479 Nil -> case f Nothing of
484 {--------------------------------------------------------------------
486 --------------------------------------------------------------------}
487 -- | The union of a list of maps.
488 unions :: [IntMap a] -> IntMap a
490 = foldlStrict union empty xs
492 -- | The union of a list of maps, with a combining operation
493 unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a
495 = foldlStrict (unionWith f) empty ts
497 -- | /O(n+m)/. The (left-biased) union of two maps.
498 -- It prefers the first map when duplicate keys are encountered,
499 -- i.e. (@'union' == 'unionWith' 'const'@).
500 union :: IntMap a -> IntMap a -> IntMap a
501 union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
502 | shorter m1 m2 = union1
503 | shorter m2 m1 = union2
504 | p1 == p2 = Bin p1 m1 (union l1 l2) (union r1 r2)
505 | otherwise = join p1 t1 p2 t2
507 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
508 | zero p2 m1 = Bin p1 m1 (union l1 t2) r1
509 | otherwise = Bin p1 m1 l1 (union r1 t2)
511 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
512 | zero p1 m2 = Bin p2 m2 (union t1 l2) r2
513 | otherwise = Bin p2 m2 l2 (union t1 r2)
515 union (Tip k x) t = insert k x t
516 union t (Tip k x) = insertWith (\x y -> y) k x t -- right bias
520 -- | /O(n+m)/. The union with a combining function.
521 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
523 = unionWithKey (\k x y -> f x y) m1 m2
525 -- | /O(n+m)/. The union with a combining function.
526 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
527 unionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
528 | shorter m1 m2 = union1
529 | shorter m2 m1 = union2
530 | p1 == p2 = Bin p1 m1 (unionWithKey f l1 l2) (unionWithKey f r1 r2)
531 | otherwise = join p1 t1 p2 t2
533 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
534 | zero p2 m1 = Bin p1 m1 (unionWithKey f l1 t2) r1
535 | otherwise = Bin p1 m1 l1 (unionWithKey f r1 t2)
537 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
538 | zero p1 m2 = Bin p2 m2 (unionWithKey f t1 l2) r2
539 | otherwise = Bin p2 m2 l2 (unionWithKey f t1 r2)
541 unionWithKey f (Tip k x) t = insertWithKey f k x t
542 unionWithKey f t (Tip k x) = insertWithKey (\k x y -> f k y x) k x t -- right bias
543 unionWithKey f Nil t = t
544 unionWithKey f t Nil = t
546 {--------------------------------------------------------------------
548 --------------------------------------------------------------------}
549 -- | /O(n+m)/. Difference between two maps (based on keys).
550 difference :: IntMap a -> IntMap b -> IntMap a
551 difference t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
552 | shorter m1 m2 = difference1
553 | shorter m2 m1 = difference2
554 | p1 == p2 = bin p1 m1 (difference l1 l2) (difference r1 r2)
557 difference1 | nomatch p2 p1 m1 = t1
558 | zero p2 m1 = bin p1 m1 (difference l1 t2) r1
559 | otherwise = bin p1 m1 l1 (difference r1 t2)
561 difference2 | nomatch p1 p2 m2 = t1
562 | zero p1 m2 = difference t1 l2
563 | otherwise = difference t1 r2
565 difference t1@(Tip k x) t2
569 difference Nil t = Nil
570 difference t (Tip k x) = delete k t
573 -- | /O(n+m)/. Difference with a combining function.
574 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
575 differenceWith f m1 m2
576 = differenceWithKey (\k x y -> f x y) m1 m2
578 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
579 -- encountered, the combining function is applied to the key and both values.
580 -- If it returns 'Nothing', the element is discarded (proper set difference).
581 -- If it returns (@'Just' y@), the element is updated with a new value @y@.
582 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
583 differenceWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
584 | shorter m1 m2 = difference1
585 | shorter m2 m1 = difference2
586 | p1 == p2 = bin p1 m1 (differenceWithKey f l1 l2) (differenceWithKey f r1 r2)
589 difference1 | nomatch p2 p1 m1 = t1
590 | zero p2 m1 = bin p1 m1 (differenceWithKey f l1 t2) r1
591 | otherwise = bin p1 m1 l1 (differenceWithKey f r1 t2)
593 difference2 | nomatch p1 p2 m2 = t1
594 | zero p1 m2 = differenceWithKey f t1 l2
595 | otherwise = differenceWithKey f t1 r2
597 differenceWithKey f t1@(Tip k x) t2
598 = case lookup k t2 of
599 Just y -> case f k x y of
604 differenceWithKey f Nil t = Nil
605 differenceWithKey f t (Tip k y) = updateWithKey (\k x -> f k x y) k t
606 differenceWithKey f t Nil = t
609 {--------------------------------------------------------------------
611 --------------------------------------------------------------------}
612 -- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys).
613 intersection :: IntMap a -> IntMap b -> IntMap a
614 intersection t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
615 | shorter m1 m2 = intersection1
616 | shorter m2 m1 = intersection2
617 | p1 == p2 = bin p1 m1 (intersection l1 l2) (intersection r1 r2)
620 intersection1 | nomatch p2 p1 m1 = Nil
621 | zero p2 m1 = intersection l1 t2
622 | otherwise = intersection r1 t2
624 intersection2 | nomatch p1 p2 m2 = Nil
625 | zero p1 m2 = intersection t1 l2
626 | otherwise = intersection t1 r2
628 intersection t1@(Tip k x) t2
631 intersection t (Tip k x)
635 intersection Nil t = Nil
636 intersection t Nil = Nil
638 -- | /O(n+m)/. The intersection with a combining function.
639 intersectionWith :: (a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
640 intersectionWith f m1 m2
641 = intersectionWithKey (\k x y -> f x y) m1 m2
643 -- | /O(n+m)/. The intersection with a combining function.
644 intersectionWithKey :: (Key -> a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
645 intersectionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
646 | shorter m1 m2 = intersection1
647 | shorter m2 m1 = intersection2
648 | p1 == p2 = bin p1 m1 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2)
651 intersection1 | nomatch p2 p1 m1 = Nil
652 | zero p2 m1 = intersectionWithKey f l1 t2
653 | otherwise = intersectionWithKey f r1 t2
655 intersection2 | nomatch p1 p2 m2 = Nil
656 | zero p1 m2 = intersectionWithKey f t1 l2
657 | otherwise = intersectionWithKey f t1 r2
659 intersectionWithKey f t1@(Tip k x) t2
660 = case lookup k t2 of
661 Just y -> Tip k (f k x y)
663 intersectionWithKey f t1 (Tip k y)
664 = case lookup k t1 of
665 Just x -> Tip k (f k x y)
667 intersectionWithKey f Nil t = Nil
668 intersectionWithKey f t Nil = Nil
671 {--------------------------------------------------------------------
673 --------------------------------------------------------------------}
674 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
675 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
676 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
677 isProperSubmapOf m1 m2
678 = isProperSubmapOfBy (==) m1 m2
680 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
681 The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
682 @m1@ and @m2@ are not equal,
683 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
684 applied to their respective values. For example, the following
685 expressions are all 'True':
687 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
688 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
690 But the following are all 'False':
692 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
693 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
694 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
696 isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
697 isProperSubmapOfBy pred t1 t2
698 = case submapCmp pred t1 t2 of
702 submapCmp pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
704 | shorter m2 m1 = submapCmpLt
705 | p1 == p2 = submapCmpEq
706 | otherwise = GT -- disjoint
708 submapCmpLt | nomatch p1 p2 m2 = GT
709 | zero p1 m2 = submapCmp pred t1 l2
710 | otherwise = submapCmp pred t1 r2
711 submapCmpEq = case (submapCmp pred l1 l2, submapCmp pred r1 r2) of
717 submapCmp pred (Bin p m l r) t = GT
718 submapCmp pred (Tip kx x) (Tip ky y)
719 | (kx == ky) && pred x y = EQ
720 | otherwise = GT -- disjoint
721 submapCmp pred (Tip k x) t
723 Just y | pred x y -> LT
724 other -> GT -- disjoint
725 submapCmp pred Nil Nil = EQ
726 submapCmp pred Nil t = LT
728 -- | /O(n+m)/. Is this a submap?
729 -- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
730 isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
732 = isSubmapOfBy (==) m1 m2
735 The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if
736 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
737 applied to their respective values. For example, the following
738 expressions are all 'True':
740 > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
741 > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
742 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
744 But the following are all 'False':
746 > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)])
747 > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
748 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
751 isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
752 isSubmapOfBy pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
753 | shorter m1 m2 = False
754 | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy pred t1 l2
755 else isSubmapOfBy pred t1 r2)
756 | otherwise = (p1==p2) && isSubmapOfBy pred l1 l2 && isSubmapOfBy pred r1 r2
757 isSubmapOfBy pred (Bin p m l r) t = False
758 isSubmapOfBy pred (Tip k x) t = case lookup k t of
761 isSubmapOfBy pred Nil t = True
763 {--------------------------------------------------------------------
765 --------------------------------------------------------------------}
766 -- | /O(n)/. Map a function over all values in the map.
767 map :: (a -> b) -> IntMap a -> IntMap b
769 = mapWithKey (\k x -> f x) m
771 -- | /O(n)/. Map a function over all values in the map.
772 mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b
775 Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r)
776 Tip k x -> Tip k (f k x)
779 -- | /O(n)/. The function @'mapAccum'@ threads an accumulating
780 -- argument through the map in ascending order of keys.
781 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
783 = mapAccumWithKey (\a k x -> f a x) a m
785 -- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating
786 -- argument through the map in ascending order of keys.
787 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
788 mapAccumWithKey f a t
791 -- | /O(n)/. The function @'mapAccumL'@ threads an accumulating
792 -- argument through the map in ascending order of keys.
793 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
796 Bin p m l r -> let (a1,l') = mapAccumL f a l
797 (a2,r') = mapAccumL f a1 r
798 in (a2,Bin p m l' r')
799 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
803 -- | /O(n)/. The function @'mapAccumR'@ threads an accumulating
804 -- argument throught the map in descending order of keys.
805 mapAccumR :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
808 Bin p m l r -> let (a1,r') = mapAccumR f a r
809 (a2,l') = mapAccumR f a1 l
810 in (a2,Bin p m l' r')
811 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
814 {--------------------------------------------------------------------
816 --------------------------------------------------------------------}
817 -- | /O(n)/. Filter all values that satisfy some predicate.
818 filter :: (a -> Bool) -> IntMap a -> IntMap a
820 = filterWithKey (\k x -> p x) m
822 -- | /O(n)/. Filter all keys\/values that satisfy some predicate.
823 filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a
827 -> bin p m (filterWithKey pred l) (filterWithKey pred r)
833 -- | /O(n)/. partition the map according to some predicate. The first
834 -- map contains all elements that satisfy the predicate, the second all
835 -- elements that fail the predicate. See also 'split'.
836 partition :: (a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
838 = partitionWithKey (\k x -> p x) m
840 -- | /O(n)/. partition the map according to some predicate. The first
841 -- map contains all elements that satisfy the predicate, the second all
842 -- elements that fail the predicate. See also 'split'.
843 partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
844 partitionWithKey pred t
847 -> let (l1,l2) = partitionWithKey pred l
848 (r1,r2) = partitionWithKey pred r
849 in (bin p m l1 r1, bin p m l2 r2)
851 | pred k x -> (t,Nil)
852 | otherwise -> (Nil,t)
855 -- | /O(n)/. Map values and collect the 'Just' results.
856 mapMaybe :: (a -> Maybe b) -> IntMap a -> IntMap b
858 = mapMaybeWithKey (\k x -> f x) m
860 -- | /O(n)/. Map keys\/values and collect the 'Just' results.
861 mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b
862 mapMaybeWithKey f (Bin p m l r)
863 = bin p m (mapMaybeWithKey f l) (mapMaybeWithKey f r)
864 mapMaybeWithKey f (Tip k x) = case f k x of
867 mapMaybeWithKey f Nil = Nil
869 -- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
870 mapEither :: (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
872 = mapEitherWithKey (\k x -> f x) m
874 -- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
875 mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
876 mapEitherWithKey f (Bin p m l r)
877 = (bin p m l1 r1, bin p m l2 r2)
879 (l1,l2) = mapEitherWithKey f l
880 (r1,r2) = mapEitherWithKey f r
881 mapEitherWithKey f (Tip k x) = case f k x of
882 Left y -> (Tip k y, Nil)
883 Right z -> (Nil, Tip k z)
884 mapEitherWithKey f Nil = (Nil, Nil)
886 -- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@
887 -- where all keys in @map1@ are lower than @k@ and all keys in
888 -- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
889 split :: Key -> IntMap a -> (IntMap a,IntMap a)
893 | m < 0 -> (if k >= 0 -- handle negative numbers.
894 then let (lt,gt) = split' k l in (union r lt, gt)
895 else let (lt,gt) = split' k r in (lt, union gt l))
896 | otherwise -> split' k t
900 | otherwise -> (Nil,Nil)
903 split' :: Key -> IntMap a -> (IntMap a,IntMap a)
907 | nomatch k p m -> if k>p then (t,Nil) else (Nil,t)
908 | zero k m -> let (lt,gt) = split k l in (lt,union gt r)
909 | otherwise -> let (lt,gt) = split k r in (union l lt,gt)
913 | otherwise -> (Nil,Nil)
916 -- | /O(log n)/. Performs a 'split' but also returns whether the pivot
917 -- key was found in the original map.
918 splitLookup :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
922 | m < 0 -> (if k >= 0 -- handle negative numbers.
923 then let (lt,found,gt) = splitLookup' k l in (union r lt,found, gt)
924 else let (lt,found,gt) = splitLookup' k r in (lt,found, union gt l))
925 | otherwise -> splitLookup' k t
927 | k>ky -> (t,Nothing,Nil)
928 | k<ky -> (Nil,Nothing,t)
929 | otherwise -> (Nil,Just y,Nil)
930 Nil -> (Nil,Nothing,Nil)
932 splitLookup' :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
936 | nomatch k p m -> if k>p then (t,Nothing,Nil) else (Nil,Nothing,t)
937 | zero k m -> let (lt,found,gt) = splitLookup k l in (lt,found,union gt r)
938 | otherwise -> let (lt,found,gt) = splitLookup k r in (union l lt,found,gt)
940 | k>ky -> (t,Nothing,Nil)
941 | k<ky -> (Nil,Nothing,t)
942 | otherwise -> (Nil,Just y,Nil)
943 Nil -> (Nil,Nothing,Nil)
945 {--------------------------------------------------------------------
947 --------------------------------------------------------------------}
948 -- | /O(n)/. Fold the values in the map, such that
949 -- @'fold' f z == 'Prelude.foldr' f z . 'elems'@.
952 -- > elems map = fold (:) [] map
954 fold :: (a -> b -> b) -> b -> IntMap a -> b
956 = foldWithKey (\k x y -> f x y) z t
958 -- | /O(n)/. Fold the keys and values in the map, such that
959 -- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
962 -- > keys map = foldWithKey (\k x ks -> k:ks) [] map
964 foldWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
968 foldr :: (Key -> a -> b -> b) -> b -> IntMap a -> b
971 Bin 0 m l r | m < 0 -> foldr' f (foldr' f z l) r -- put negative numbers before.
972 Bin _ _ _ _ -> foldr' f z t
976 foldr' :: (Key -> a -> b -> b) -> b -> IntMap a -> b
979 Bin p m l r -> foldr' f (foldr' f z r) l
985 {--------------------------------------------------------------------
987 --------------------------------------------------------------------}
989 -- Return all elements of the map in the ascending order of their keys.
990 elems :: IntMap a -> [a]
992 = foldWithKey (\k x xs -> x:xs) [] m
994 -- | /O(n)/. Return all keys of the map in ascending order.
995 keys :: IntMap a -> [Key]
997 = foldWithKey (\k x ks -> k:ks) [] m
999 -- | /O(n*min(n,W))/. The set of all keys of the map.
1000 keysSet :: IntMap a -> IntSet.IntSet
1001 keysSet m = IntSet.fromDistinctAscList (keys m)
1004 -- | /O(n)/. Return all key\/value pairs in the map in ascending key order.
1005 assocs :: IntMap a -> [(Key,a)]
1010 {--------------------------------------------------------------------
1012 --------------------------------------------------------------------}
1013 -- | /O(n)/. Convert the map to a list of key\/value pairs.
1014 toList :: IntMap a -> [(Key,a)]
1016 = foldWithKey (\k x xs -> (k,x):xs) [] t
1018 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
1019 -- keys are in ascending order.
1020 toAscList :: IntMap a -> [(Key,a)]
1022 = -- NOTE: the following algorithm only works for big-endian trees
1023 let (pos,neg) = span (\(k,x) -> k >=0) (foldr (\k x xs -> (k,x):xs) [] t) in neg ++ pos
1025 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
1026 fromList :: [(Key,a)] -> IntMap a
1028 = foldlStrict ins empty xs
1030 ins t (k,x) = insert k x t
1032 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
1033 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1035 = fromListWithKey (\k x y -> f x y) xs
1037 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
1038 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1039 fromListWithKey f xs
1040 = foldlStrict ins empty xs
1042 ins t (k,x) = insertWithKey f k x t
1044 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
1045 -- the keys are in ascending order.
1046 fromAscList :: [(Key,a)] -> IntMap a
1050 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
1051 -- the keys are in ascending order, with a combining function on equal keys.
1052 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1053 fromAscListWith f xs
1056 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
1057 -- the keys are in ascending order, with a combining function on equal keys.
1058 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1059 fromAscListWithKey f xs
1060 = fromListWithKey f xs
1062 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
1063 -- the keys are in ascending order and all distinct.
1064 fromDistinctAscList :: [(Key,a)] -> IntMap a
1065 fromDistinctAscList xs
1069 {--------------------------------------------------------------------
1071 --------------------------------------------------------------------}
1072 instance Eq a => Eq (IntMap a) where
1073 t1 == t2 = equal t1 t2
1074 t1 /= t2 = nequal t1 t2
1076 equal :: Eq a => IntMap a -> IntMap a -> Bool
1077 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1078 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
1079 equal (Tip kx x) (Tip ky y)
1080 = (kx == ky) && (x==y)
1081 equal Nil Nil = True
1084 nequal :: Eq a => IntMap a -> IntMap a -> Bool
1085 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1086 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
1087 nequal (Tip kx x) (Tip ky y)
1088 = (kx /= ky) || (x/=y)
1089 nequal Nil Nil = False
1092 {--------------------------------------------------------------------
1094 --------------------------------------------------------------------}
1096 instance Ord a => Ord (IntMap a) where
1097 compare m1 m2 = compare (toList m1) (toList m2)
1099 {--------------------------------------------------------------------
1101 --------------------------------------------------------------------}
1103 instance Functor IntMap where
1106 {--------------------------------------------------------------------
1108 --------------------------------------------------------------------}
1110 instance Show a => Show (IntMap a) where
1111 showsPrec d m = showParen (d > 10) $
1112 showString "fromList " . shows (toList m)
1114 showMap :: (Show a) => [(Key,a)] -> ShowS
1118 = showChar '{' . showElem x . showTail xs
1120 showTail [] = showChar '}'
1121 showTail (x:xs) = showChar ',' . showElem x . showTail xs
1123 showElem (k,x) = shows k . showString ":=" . shows x
1125 {--------------------------------------------------------------------
1127 --------------------------------------------------------------------}
1128 instance (Read e) => Read (IntMap e) where
1129 #ifdef __GLASGOW_HASKELL__
1130 readPrec = parens $ prec 10 $ do
1131 Ident "fromList" <- lexP
1133 return (fromList xs)
1135 readListPrec = readListPrecDefault
1137 readsPrec p = readParen (p > 10) $ \ r -> do
1138 ("fromList",s) <- lex r
1140 return (fromList xs,t)
1143 {--------------------------------------------------------------------
1145 --------------------------------------------------------------------}
1147 #include "Typeable.h"
1148 INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap")
1150 {--------------------------------------------------------------------
1152 --------------------------------------------------------------------}
1153 -- | /O(n)/. Show the tree that implements the map. The tree is shown
1154 -- in a compressed, hanging format.
1155 showTree :: Show a => IntMap a -> String
1157 = showTreeWith True False s
1160 {- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows
1161 the tree that implements the map. If @hang@ is
1162 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
1163 @wide@ is 'True', an extra wide version is shown.
1165 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
1166 showTreeWith hang wide t
1167 | hang = (showsTreeHang wide [] t) ""
1168 | otherwise = (showsTree wide [] [] t) ""
1170 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
1171 showsTree wide lbars rbars t
1174 -> showsTree wide (withBar rbars) (withEmpty rbars) r .
1175 showWide wide rbars .
1176 showsBars lbars . showString (showBin p m) . showString "\n" .
1177 showWide wide lbars .
1178 showsTree wide (withEmpty lbars) (withBar lbars) l
1180 -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1181 Nil -> showsBars lbars . showString "|\n"
1183 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
1184 showsTreeHang wide bars t
1187 -> showsBars bars . showString (showBin p m) . showString "\n" .
1188 showWide wide bars .
1189 showsTreeHang wide (withBar bars) l .
1190 showWide wide bars .
1191 showsTreeHang wide (withEmpty bars) r
1193 -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1194 Nil -> showsBars bars . showString "|\n"
1197 = "*" -- ++ show (p,m)
1200 | wide = showString (concat (reverse bars)) . showString "|\n"
1203 showsBars :: [String] -> ShowS
1207 _ -> showString (concat (reverse (tail bars))) . showString node
1210 withBar bars = "| ":bars
1211 withEmpty bars = " ":bars
1214 {--------------------------------------------------------------------
1216 --------------------------------------------------------------------}
1217 {--------------------------------------------------------------------
1219 --------------------------------------------------------------------}
1220 join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
1222 | zero p1 m = Bin p m t1 t2
1223 | otherwise = Bin p m t2 t1
1225 m = branchMask p1 p2
1228 {--------------------------------------------------------------------
1229 @bin@ assures that we never have empty trees within a tree.
1230 --------------------------------------------------------------------}
1231 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
1234 bin p m l r = Bin p m l r
1237 {--------------------------------------------------------------------
1238 Endian independent bit twiddling
1239 --------------------------------------------------------------------}
1240 zero :: Key -> Mask -> Bool
1242 = (natFromInt i) .&. (natFromInt m) == 0
1244 nomatch,match :: Key -> Prefix -> Mask -> Bool
1251 mask :: Key -> Mask -> Prefix
1253 = maskW (natFromInt i) (natFromInt m)
1256 zeroN :: Nat -> Nat -> Bool
1257 zeroN i m = (i .&. m) == 0
1259 {--------------------------------------------------------------------
1260 Big endian operations
1261 --------------------------------------------------------------------}
1262 maskW :: Nat -> Nat -> Prefix
1264 = intFromNat (i .&. (complement (m-1) `xor` m))
1266 shorter :: Mask -> Mask -> Bool
1268 = (natFromInt m1) > (natFromInt m2)
1270 branchMask :: Prefix -> Prefix -> Mask
1272 = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
1274 {----------------------------------------------------------------------
1275 Finding the highest bit (mask) in a word [x] can be done efficiently in
1277 * convert to a floating point value and the mantissa tells us the
1278 [log2(x)] that corresponds with the highest bit position. The mantissa
1279 is retrieved either via the standard C function [frexp] or by some bit
1280 twiddling on IEEE compatible numbers (float). Note that one needs to
1281 use at least [double] precision for an accurate mantissa of 32 bit
1283 * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
1284 * use processor specific assembler instruction (asm).
1286 The most portable way would be [bit], but is it efficient enough?
1287 I have measured the cycle counts of the different methods on an AMD
1288 Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
1290 highestBitMask: method cycles
1297 highestBit: method cycles
1304 Wow, the bit twiddling is on today's RISC like machines even faster
1305 than a single CISC instruction (BSR)!
1306 ----------------------------------------------------------------------}
1308 {----------------------------------------------------------------------
1309 [highestBitMask] returns a word where only the highest bit is set.
1310 It is found by first setting all bits in lower positions than the
1311 highest bit and than taking an exclusive or with the original value.
1312 Allthough the function may look expensive, GHC compiles this into
1313 excellent C code that subsequently compiled into highly efficient
1314 machine code. The algorithm is derived from Jorg Arndt's FXT library.
1315 ----------------------------------------------------------------------}
1316 highestBitMask :: Nat -> Nat
1318 = case (x .|. shiftRL x 1) of
1319 x -> case (x .|. shiftRL x 2) of
1320 x -> case (x .|. shiftRL x 4) of
1321 x -> case (x .|. shiftRL x 8) of
1322 x -> case (x .|. shiftRL x 16) of
1323 x -> case (x .|. shiftRL x 32) of -- for 64 bit platforms
1324 x -> (x `xor` (shiftRL x 1))
1327 {--------------------------------------------------------------------
1329 --------------------------------------------------------------------}
1333 (x:xx) -> let z' = f z x in seq z' (foldlStrict f z' xx)
1336 {--------------------------------------------------------------------
1338 --------------------------------------------------------------------}
1339 testTree :: [Int] -> IntMap Int
1340 testTree xs = fromList [(x,x*x*30696 `mod` 65521) | x <- xs]
1341 test1 = testTree [1..20]
1342 test2 = testTree [30,29..10]
1343 test3 = testTree [1,4,6,89,2323,53,43,234,5,79,12,9,24,9,8,423,8,42,4,8,9,3]
1345 {--------------------------------------------------------------------
1347 --------------------------------------------------------------------}
1352 { configMaxTest = 500
1353 , configMaxFail = 5000
1354 , configSize = \n -> (div n 2 + 3)
1355 , configEvery = \n args -> let s = show n in s ++ [ '\b' | _ <- s ]
1359 {--------------------------------------------------------------------
1360 Arbitrary, reasonably balanced trees
1361 --------------------------------------------------------------------}
1362 instance Arbitrary a => Arbitrary (IntMap a) where
1363 arbitrary = do{ ks <- arbitrary
1364 ; xs <- mapM (\k -> do{ x <- arbitrary; return (k,x)}) ks
1365 ; return (fromList xs)
1369 {--------------------------------------------------------------------
1370 Single, Insert, Delete
1371 --------------------------------------------------------------------}
1372 prop_Single :: Key -> Int -> Bool
1374 = (insert k x empty == singleton k x)
1376 prop_InsertDelete :: Key -> Int -> IntMap Int -> Property
1377 prop_InsertDelete k x t
1378 = not (member k t) ==> delete k (insert k x t) == t
1380 prop_UpdateDelete :: Key -> IntMap Int -> Bool
1381 prop_UpdateDelete k t
1382 = update (const Nothing) k t == delete k t
1385 {--------------------------------------------------------------------
1387 --------------------------------------------------------------------}
1388 prop_UnionInsert :: Key -> Int -> IntMap Int -> Bool
1389 prop_UnionInsert k x t
1390 = union (singleton k x) t == insert k x t
1392 prop_UnionAssoc :: IntMap Int -> IntMap Int -> IntMap Int -> Bool
1393 prop_UnionAssoc t1 t2 t3
1394 = union t1 (union t2 t3) == union (union t1 t2) t3
1396 prop_UnionComm :: IntMap Int -> IntMap Int -> Bool
1397 prop_UnionComm t1 t2
1398 = (union t1 t2 == unionWith (\x y -> y) t2 t1)
1401 prop_Diff :: [(Key,Int)] -> [(Key,Int)] -> Bool
1403 = List.sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys)))
1404 == List.sort ((List.\\) (nub (Prelude.map fst xs)) (nub (Prelude.map fst ys)))
1406 prop_Int :: [(Key,Int)] -> [(Key,Int)] -> Bool
1408 = List.sort (keys (intersection (fromListWith (+) xs) (fromListWith (+) ys)))
1409 == List.sort (nub ((List.intersect) (Prelude.map fst xs) (Prelude.map fst ys)))
1411 {--------------------------------------------------------------------
1413 --------------------------------------------------------------------}
1415 = forAll (choose (5,100)) $ \n ->
1416 let xs = [(x,()) | x <- [0..n::Int]]
1417 in fromAscList xs == fromList xs
1419 prop_List :: [Key] -> Bool
1421 = (sort (nub xs) == [x | (x,()) <- toAscList (fromList [(x,()) | x <- xs])])