1 {-# OPTIONS -cpp -fglasgow-exts #-}
2 --------------------------------------------------------------------------------
3 {-| Module : Data.IntMap
4 Copyright : (c) Daan Leijen 2002
6 Maintainer : libraries@haskell.org
7 Stability : provisional
10 An efficient implementation of maps from integer keys to values.
12 This module is intended to be imported @qualified@, to avoid name
13 clashes with Prelude functions. eg.
15 > import Data.IntMap as Map
17 The implementation is based on /big-endian patricia trees/. This data structure
18 performs especially well on binary operations like 'union' and 'intersection'. However,
19 my benchmarks show that it is also (much) faster on insertions and deletions when
20 compared to a generic size-balanced map implementation (see "Map" and "Data.FiniteMap").
22 * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\",
23 Workshop on ML, September 1998, pages 77--86, <http://www.cse.ogi.edu/~andy/pub/finite.htm>
25 * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve Information
26 Coded In Alphanumeric/\", Journal of the ACM, 15(4), October 1968, pages 514--534.
28 Many operations have a worst-case complexity of /O(min(n,W))/. This means that the
29 operation can become linear in the number of elements
30 with a maximum of /W/ -- the number of bits in an 'Int' (32 or 64).
32 ---------------------------------------------------------------------------------
35 IntMap, Key -- instance Eq,Show
53 , insertWith, insertWithKey, insertLookupWithKey
110 , fromDistinctAscList
122 , isSubmapOf, isSubmapOfBy
123 , isProperSubmapOf, isProperSubmapOfBy
131 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
135 import qualified Data.IntSet as IntSet
139 import qualified Prelude
140 import Debug.QuickCheck
141 import List (nub,sort)
142 import qualified List
145 #ifdef __GLASGOW_HASKELL__
146 {--------------------------------------------------------------------
147 GHC: use unboxing to get @shiftRL@ inlined.
148 --------------------------------------------------------------------}
149 #if __GLASGOW_HASKELL__ >= 503
151 import GHC.Exts ( Word(..), Int(..), shiftRL# )
154 import GlaExts ( Word(..), Int(..), shiftRL# )
157 infixl 9 \\{-This comment teaches CPP correct behaviour -}
161 natFromInt :: Key -> Nat
162 natFromInt i = fromIntegral i
164 intFromNat :: Nat -> Key
165 intFromNat w = fromIntegral w
167 shiftRL :: Nat -> Key -> Nat
168 shiftRL (W# x) (I# i)
172 {--------------------------------------------------------------------
174 * raises errors on boundary values when using 'fromIntegral'
175 but not with the deprecated 'fromInt/toInt'.
176 * Older Hugs doesn't define 'Word'.
177 * Newer Hugs defines 'Word' in the Prelude but no operations.
178 --------------------------------------------------------------------}
182 type Nat = Word32 -- illegal on 64-bit platforms!
184 natFromInt :: Key -> Nat
185 natFromInt i = fromInt i
187 intFromNat :: Nat -> Key
188 intFromNat w = toInt w
190 shiftRL :: Nat -> Key -> Nat
191 shiftRL x i = shiftR x i
194 {--------------------------------------------------------------------
196 * A "Nat" is a natural machine word (an unsigned Int)
197 --------------------------------------------------------------------}
203 natFromInt :: Key -> Nat
204 natFromInt i = fromIntegral i
206 intFromNat :: Nat -> Key
207 intFromNat w = fromIntegral w
209 shiftRL :: Nat -> Key -> Nat
210 shiftRL w i = shiftR w i
215 {--------------------------------------------------------------------
217 --------------------------------------------------------------------}
219 -- | /O(min(n,W))/. Find the value of a key. Calls @error@ when the element can not be found.
221 (!) :: IntMap a -> Key -> a
224 -- | /O(n+m)/. See 'difference'.
225 (\\) :: IntMap a -> IntMap b -> IntMap a
226 m1 \\ m2 = difference m1 m2
228 {--------------------------------------------------------------------
230 --------------------------------------------------------------------}
231 -- | A map of integers to values @a@.
233 | Tip {-# UNPACK #-} !Key a
234 | Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(IntMap a) !(IntMap a)
240 {--------------------------------------------------------------------
242 --------------------------------------------------------------------}
243 -- | /O(1)/. Is the map empty?
244 null :: IntMap a -> Bool
248 -- | /O(n)/. Number of elements in the map.
249 size :: IntMap a -> Int
252 Bin p m l r -> size l + size r
256 -- | /O(min(n,W))/. Is the key a member of the map?
257 member :: Key -> IntMap a -> Bool
263 -- | /O(min(n,W))/. Lookup the value of a key in the map.
264 lookup :: Key -> IntMap a -> Maybe a
266 = let nk = natFromInt k in seq nk (lookupN nk t)
268 lookupN :: Nat -> IntMap a -> Maybe a
272 | zeroN k (natFromInt m) -> lookupN k l
273 | otherwise -> lookupN k r
275 | (k == natFromInt kx) -> Just x
276 | otherwise -> Nothing
279 find' :: Key -> IntMap a -> a
282 Nothing -> error ("IntMap.find: key " ++ show k ++ " is not an element of the map")
286 -- | /O(min(n,W))/. The expression @(findWithDefault def k map)@ returns the value of key @k@ or returns @def@ when
287 -- the key is not an element of the map.
288 findWithDefault :: a -> Key -> IntMap a -> a
289 findWithDefault def k m
294 {--------------------------------------------------------------------
296 --------------------------------------------------------------------}
297 -- | /O(1)/. The empty map.
302 -- | /O(1)/. A map of one element.
303 singleton :: Key -> a -> IntMap a
307 {--------------------------------------------------------------------
309 'insert' is the inlined version of 'insertWith (\k x y -> x)'
310 --------------------------------------------------------------------}
311 -- | /O(min(n,W))/. Insert a new key\/value pair in the map. When the key
312 -- is already an element of the set, its value is replaced by the new value,
313 -- ie. 'insert' is left-biased.
314 insert :: Key -> a -> IntMap a -> IntMap a
318 | nomatch k p m -> join k (Tip k x) p t
319 | zero k m -> Bin p m (insert k x l) r
320 | otherwise -> Bin p m l (insert k x r)
323 | otherwise -> join k (Tip k x) ky t
326 -- right-biased insertion, used by 'union'
327 -- | /O(min(n,W))/. Insert with a combining function.
328 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
330 = insertWithKey (\k x y -> f x y) k x t
332 -- | /O(min(n,W))/. Insert with a combining function.
333 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
334 insertWithKey f k x t
337 | nomatch k p m -> join k (Tip k x) p t
338 | zero k m -> Bin p m (insertWithKey f k x l) r
339 | otherwise -> Bin p m l (insertWithKey f k x r)
341 | k==ky -> Tip k (f k x y)
342 | otherwise -> join k (Tip k x) ky t
346 -- | /O(min(n,W))/. The expression (@insertLookupWithKey f k x map@) is a pair where
347 -- the first element is equal to (@lookup k map@) and the second element
348 -- equal to (@insertWithKey f k x map@).
349 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
350 insertLookupWithKey f k x t
353 | nomatch k p m -> (Nothing,join k (Tip k x) p t)
354 | zero k m -> let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r)
355 | otherwise -> let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r')
357 | k==ky -> (Just y,Tip k (f k x y))
358 | otherwise -> (Nothing,join k (Tip k x) ky t)
359 Nil -> (Nothing,Tip k x)
362 {--------------------------------------------------------------------
364 [delete] is the inlined version of [deleteWith (\k x -> Nothing)]
365 --------------------------------------------------------------------}
366 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
367 -- a member of the map, the original map is returned.
368 delete :: Key -> IntMap a -> IntMap a
373 | zero k m -> bin p m (delete k l) r
374 | otherwise -> bin p m l (delete k r)
380 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
381 -- a member of the map, the original map is returned.
382 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
384 = adjustWithKey (\k x -> f x) k m
386 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
387 -- a member of the map, the original map is returned.
388 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
390 = updateWithKey (\k x -> Just (f k x)) k m
392 -- | /O(min(n,W))/. The expression (@update f k map@) updates the value @x@
393 -- at @k@ (if it is in the map). If (@f x@) is @Nothing@, the element is
394 -- deleted. If it is (@Just y@), the key @k@ is bound to the new value @y@.
395 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
397 = updateWithKey (\k x -> f x) k m
399 -- | /O(min(n,W))/. The expression (@update f k map@) updates the value @x@
400 -- at @k@ (if it is in the map). If (@f k x@) is @Nothing@, the element is
401 -- deleted. If it is (@Just y@), the key @k@ is bound to the new value @y@.
402 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
407 | zero k m -> bin p m (updateWithKey f k l) r
408 | otherwise -> bin p m l (updateWithKey f k r)
410 | k==ky -> case (f k y) of
416 -- | /O(min(n,W))/. Lookup and update.
417 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
418 updateLookupWithKey f k t
421 | nomatch k p m -> (Nothing,t)
422 | zero k m -> let (found,l') = updateLookupWithKey f k l in (found,bin p m l' r)
423 | otherwise -> let (found,r') = updateLookupWithKey f k r in (found,bin p m l r')
425 | k==ky -> case (f k y) of
426 Just y' -> (Just y,Tip ky y')
427 Nothing -> (Just y,Nil)
428 | otherwise -> (Nothing,t)
432 {--------------------------------------------------------------------
434 --------------------------------------------------------------------}
435 -- | The union of a list of maps.
436 unions :: [IntMap a] -> IntMap a
438 = foldlStrict union empty xs
440 -- | The union of a list of maps, with a combining operation
441 unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a
443 = foldlStrict (unionWith f) empty ts
445 -- | /O(n+m)/. The (left-biased) union of two sets.
446 union :: IntMap a -> IntMap a -> IntMap a
447 union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
448 | shorter m1 m2 = union1
449 | shorter m2 m1 = union2
450 | p1 == p2 = Bin p1 m1 (union l1 l2) (union r1 r2)
451 | otherwise = join p1 t1 p2 t2
453 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
454 | zero p2 m1 = Bin p1 m1 (union l1 t2) r1
455 | otherwise = Bin p1 m1 l1 (union r1 t2)
457 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
458 | zero p1 m2 = Bin p2 m2 (union t1 l2) r2
459 | otherwise = Bin p2 m2 l2 (union t1 r2)
461 union (Tip k x) t = insert k x t
462 union t (Tip k x) = insertWith (\x y -> y) k x t -- right bias
466 -- | /O(n+m)/. The union with a combining function.
467 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
469 = unionWithKey (\k x y -> f x y) m1 m2
471 -- | /O(n+m)/. The union with a combining function.
472 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
473 unionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
474 | shorter m1 m2 = union1
475 | shorter m2 m1 = union2
476 | p1 == p2 = Bin p1 m1 (unionWithKey f l1 l2) (unionWithKey f r1 r2)
477 | otherwise = join p1 t1 p2 t2
479 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
480 | zero p2 m1 = Bin p1 m1 (unionWithKey f l1 t2) r1
481 | otherwise = Bin p1 m1 l1 (unionWithKey f r1 t2)
483 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
484 | zero p1 m2 = Bin p2 m2 (unionWithKey f t1 l2) r2
485 | otherwise = Bin p2 m2 l2 (unionWithKey f t1 r2)
487 unionWithKey f (Tip k x) t = insertWithKey f k x t
488 unionWithKey f t (Tip k x) = insertWithKey (\k x y -> f k y x) k x t -- right bias
489 unionWithKey f Nil t = t
490 unionWithKey f t Nil = t
492 {--------------------------------------------------------------------
494 --------------------------------------------------------------------}
495 -- | /O(n+m)/. Difference between two maps (based on keys).
496 difference :: IntMap a -> IntMap b -> IntMap a
497 difference t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
498 | shorter m1 m2 = difference1
499 | shorter m2 m1 = difference2
500 | p1 == p2 = bin p1 m1 (difference l1 l2) (difference r1 r2)
503 difference1 | nomatch p2 p1 m1 = t1
504 | zero p2 m1 = bin p1 m1 (difference l1 t2) r1
505 | otherwise = bin p1 m1 l1 (difference r1 t2)
507 difference2 | nomatch p1 p2 m2 = t1
508 | zero p1 m2 = difference t1 l2
509 | otherwise = difference t1 r2
511 difference t1@(Tip k x) t2
515 difference Nil t = Nil
516 difference t (Tip k x) = delete k t
519 -- | /O(n+m)/. Difference with a combining function.
520 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
521 differenceWith f m1 m2
522 = differenceWithKey (\k x y -> f x y) m1 m2
524 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
525 -- encountered, the combining function is applied to the key and both values.
526 -- If it returns @Nothing@, the element is discarded (proper set difference). If
527 -- it returns (@Just y@), the element is updated with a new value @y@.
528 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
529 differenceWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
530 | shorter m1 m2 = difference1
531 | shorter m2 m1 = difference2
532 | p1 == p2 = bin p1 m1 (differenceWithKey f l1 l2) (differenceWithKey f r1 r2)
535 difference1 | nomatch p2 p1 m1 = t1
536 | zero p2 m1 = bin p1 m1 (differenceWithKey f l1 t2) r1
537 | otherwise = bin p1 m1 l1 (differenceWithKey f r1 t2)
539 difference2 | nomatch p1 p2 m2 = t1
540 | zero p1 m2 = differenceWithKey f t1 l2
541 | otherwise = differenceWithKey f t1 r2
543 differenceWithKey f t1@(Tip k x) t2
544 = case lookup k t2 of
545 Just y -> case f k x y of
550 differenceWithKey f Nil t = Nil
551 differenceWithKey f t (Tip k y) = updateWithKey (\k x -> f k x y) k t
552 differenceWithKey f t Nil = t
555 {--------------------------------------------------------------------
557 --------------------------------------------------------------------}
558 -- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys).
559 intersection :: IntMap a -> IntMap b -> IntMap a
560 intersection t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
561 | shorter m1 m2 = intersection1
562 | shorter m2 m1 = intersection2
563 | p1 == p2 = bin p1 m1 (intersection l1 l2) (intersection r1 r2)
566 intersection1 | nomatch p2 p1 m1 = Nil
567 | zero p2 m1 = intersection l1 t2
568 | otherwise = intersection r1 t2
570 intersection2 | nomatch p1 p2 m2 = Nil
571 | zero p1 m2 = intersection t1 l2
572 | otherwise = intersection t1 r2
574 intersection t1@(Tip k x) t2
577 intersection t (Tip k x)
581 intersection Nil t = Nil
582 intersection t Nil = Nil
584 -- | /O(n+m)/. The intersection with a combining function.
585 intersectionWith :: (a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
586 intersectionWith f m1 m2
587 = intersectionWithKey (\k x y -> f x y) m1 m2
589 -- | /O(n+m)/. The intersection with a combining function.
590 intersectionWithKey :: (Key -> a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
591 intersectionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
592 | shorter m1 m2 = intersection1
593 | shorter m2 m1 = intersection2
594 | p1 == p2 = bin p1 m1 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2)
597 intersection1 | nomatch p2 p1 m1 = Nil
598 | zero p2 m1 = intersectionWithKey f l1 t2
599 | otherwise = intersectionWithKey f r1 t2
601 intersection2 | nomatch p1 p2 m2 = Nil
602 | zero p1 m2 = intersectionWithKey f t1 l2
603 | otherwise = intersectionWithKey f t1 r2
605 intersectionWithKey f t1@(Tip k x) t2
606 = case lookup k t2 of
607 Just y -> Tip k (f k x y)
609 intersectionWithKey f t1 (Tip k y)
610 = case lookup k t1 of
611 Just x -> Tip k (f k x y)
613 intersectionWithKey f Nil t = Nil
614 intersectionWithKey f t Nil = Nil
617 {--------------------------------------------------------------------
619 --------------------------------------------------------------------}
620 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
621 -- Defined as (@isProperSubmapOf = isProperSubmapOfBy (==)@).
622 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
623 isProperSubmapOf m1 m2
624 = isProperSubmapOfBy (==) m1 m2
626 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
627 The expression (@isProperSubmapOfBy f m1 m2@) returns @True@ when
628 @m1@ and @m2@ are not equal,
629 all keys in @m1@ are in @m2@, and when @f@ returns @True@ when
630 applied to their respective values. For example, the following
631 expressions are all @True@.
633 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
634 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
636 But the following are all @False@:
638 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
639 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
640 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
642 isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
643 isProperSubmapOfBy pred t1 t2
644 = case submapCmp pred t1 t2 of
648 submapCmp pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
650 | shorter m2 m1 = submapCmpLt
651 | p1 == p2 = submapCmpEq
652 | otherwise = GT -- disjoint
654 submapCmpLt | nomatch p1 p2 m2 = GT
655 | zero p1 m2 = submapCmp pred t1 l2
656 | otherwise = submapCmp pred t1 r2
657 submapCmpEq = case (submapCmp pred l1 l2, submapCmp pred r1 r2) of
663 submapCmp pred (Bin p m l r) t = GT
664 submapCmp pred (Tip kx x) (Tip ky y)
665 | (kx == ky) && pred x y = EQ
666 | otherwise = GT -- disjoint
667 submapCmp pred (Tip k x) t
669 Just y | pred x y -> LT
670 other -> GT -- disjoint
671 submapCmp pred Nil Nil = EQ
672 submapCmp pred Nil t = LT
674 -- | /O(n+m)/. Is this a submap? Defined as (@isSubmapOf = isSubmapOfBy (==)@).
675 isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
677 = isSubmapOfBy (==) m1 m2
680 The expression (@isSubmapOfBy f m1 m2@) returns @True@ if
681 all keys in @m1@ are in @m2@, and when @f@ returns @True@ when
682 applied to their respective values. For example, the following
683 expressions are all @True@.
685 > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
686 > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
687 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
689 But the following are all @False@:
691 > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)])
692 > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
693 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
696 isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
697 isSubmapOfBy pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
698 | shorter m1 m2 = False
699 | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy pred t1 l2
700 else isSubmapOfBy pred t1 r2)
701 | otherwise = (p1==p2) && isSubmapOfBy pred l1 l2 && isSubmapOfBy pred r1 r2
702 isSubmapOfBy pred (Bin p m l r) t = False
703 isSubmapOfBy pred (Tip k x) t = case lookup k t of
706 isSubmapOfBy pred Nil t = True
708 {--------------------------------------------------------------------
710 --------------------------------------------------------------------}
711 -- | /O(n)/. Map a function over all values in the map.
712 map :: (a -> b) -> IntMap a -> IntMap b
714 = mapWithKey (\k x -> f x) m
716 -- | /O(n)/. Map a function over all values in the map.
717 mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b
720 Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r)
721 Tip k x -> Tip k (f k x)
724 -- | /O(n)/. The function @mapAccum@ threads an accumulating
725 -- argument through the map in an unspecified order.
726 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
728 = mapAccumWithKey (\a k x -> f a x) a m
730 -- | /O(n)/. The function @mapAccumWithKey@ threads an accumulating
731 -- argument through the map in an unspecified order.
732 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
733 mapAccumWithKey f a t
736 -- | /O(n)/. The function @mapAccumL@ threads an accumulating
737 -- argument through the map in pre-order.
738 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
741 Bin p m l r -> let (a1,l') = mapAccumL f a l
742 (a2,r') = mapAccumL f a1 r
743 in (a2,Bin p m l' r')
744 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
748 -- | /O(n)/. The function @mapAccumR@ threads an accumulating
749 -- argument throught the map in post-order.
750 mapAccumR :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
753 Bin p m l r -> let (a1,r') = mapAccumR f a r
754 (a2,l') = mapAccumR f a1 l
755 in (a2,Bin p m l' r')
756 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
759 {--------------------------------------------------------------------
761 --------------------------------------------------------------------}
762 -- | /O(n)/. Filter all values that satisfy some predicate.
763 filter :: (a -> Bool) -> IntMap a -> IntMap a
765 = filterWithKey (\k x -> p x) m
767 -- | /O(n)/. Filter all keys\/values that satisfy some predicate.
768 filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a
772 -> bin p m (filterWithKey pred l) (filterWithKey pred r)
778 -- | /O(n)/. partition the map according to some predicate. The first
779 -- map contains all elements that satisfy the predicate, the second all
780 -- elements that fail the predicate. See also 'split'.
781 partition :: (a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
783 = partitionWithKey (\k x -> p x) m
785 -- | /O(n)/. partition the map according to some predicate. The first
786 -- map contains all elements that satisfy the predicate, the second all
787 -- elements that fail the predicate. See also 'split'.
788 partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
789 partitionWithKey pred t
792 -> let (l1,l2) = partitionWithKey pred l
793 (r1,r2) = partitionWithKey pred r
794 in (bin p m l1 r1, bin p m l2 r2)
796 | pred k x -> (t,Nil)
797 | otherwise -> (Nil,t)
801 -- | /O(log n)/. The expression (@split k map@) is a pair @(map1,map2)@
802 -- where all keys in @map1@ are lower than @k@ and all keys in
803 -- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
804 split :: Key -> IntMap a -> (IntMap a,IntMap a)
808 | zero k m -> let (lt,gt) = split k l in (lt,union gt r)
809 | otherwise -> let (lt,gt) = split k r in (union l lt,gt)
813 | otherwise -> (Nil,Nil)
816 -- | /O(log n)/. Performs a 'split' but also returns whether the pivot
817 -- key was found in the original map.
818 splitLookup :: Key -> IntMap a -> (Maybe a,IntMap a,IntMap a)
822 | zero k m -> let (found,lt,gt) = splitLookup k l in (found,lt,union gt r)
823 | otherwise -> let (found,lt,gt) = splitLookup k r in (found,union l lt,gt)
825 | k>ky -> (Nothing,t,Nil)
826 | k<ky -> (Nothing,Nil,t)
827 | otherwise -> (Just y,Nil,Nil)
828 Nil -> (Nothing,Nil,Nil)
830 {--------------------------------------------------------------------
832 --------------------------------------------------------------------}
833 -- | /O(n)/. Fold over the elements of a map in an unspecified order.
835 -- > sum map = fold (+) 0 map
836 -- > elems map = fold (:) [] map
837 fold :: (a -> b -> b) -> b -> IntMap a -> b
839 = foldWithKey (\k x y -> f x y) z t
841 -- | /O(n)/. Fold over the elements of a map in an unspecified order.
843 -- > keys map = foldWithKey (\k x ks -> k:ks) [] map
844 foldWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
848 foldr :: (Key -> a -> b -> b) -> b -> IntMap a -> b
851 Bin p m l r -> foldr f (foldr f z r) l
855 {--------------------------------------------------------------------
857 --------------------------------------------------------------------}
858 -- | /O(n)/. Return all elements of the map.
859 elems :: IntMap a -> [a]
861 = foldWithKey (\k x xs -> x:xs) [] m
863 -- | /O(n)/. Return all keys of the map.
864 keys :: IntMap a -> [Key]
866 = foldWithKey (\k x ks -> k:ks) [] m
868 -- | /O(n*min(n,W))/. The set of all keys of the map.
869 keysSet :: IntMap a -> IntSet.IntSet
870 keysSet m = IntSet.fromDistinctAscList (keys m)
873 -- | /O(n)/. Return all key\/value pairs in the map.
874 assocs :: IntMap a -> [(Key,a)]
879 {--------------------------------------------------------------------
881 --------------------------------------------------------------------}
882 -- | /O(n)/. Convert the map to a list of key\/value pairs.
883 toList :: IntMap a -> [(Key,a)]
885 = foldWithKey (\k x xs -> (k,x):xs) [] t
887 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
888 -- keys are in ascending order.
889 toAscList :: IntMap a -> [(Key,a)]
891 = -- NOTE: the following algorithm only works for big-endian trees
892 let (pos,neg) = span (\(k,x) -> k >=0) (foldr (\k x xs -> (k,x):xs) [] t) in neg ++ pos
894 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
895 fromList :: [(Key,a)] -> IntMap a
897 = foldlStrict ins empty xs
899 ins t (k,x) = insert k x t
901 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
902 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
904 = fromListWithKey (\k x y -> f x y) xs
906 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
907 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
909 = foldlStrict ins empty xs
911 ins t (k,x) = insertWithKey f k x t
913 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
914 -- the keys are in ascending order.
915 fromAscList :: [(Key,a)] -> IntMap a
919 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
920 -- the keys are in ascending order, with a combining function on equal keys.
921 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
925 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
926 -- the keys are in ascending order, with a combining function on equal keys.
927 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
928 fromAscListWithKey f xs
929 = fromListWithKey f xs
931 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
932 -- the keys are in ascending order and all distinct.
933 fromDistinctAscList :: [(Key,a)] -> IntMap a
934 fromDistinctAscList xs
938 {--------------------------------------------------------------------
940 --------------------------------------------------------------------}
941 instance Eq a => Eq (IntMap a) where
942 t1 == t2 = equal t1 t2
943 t1 /= t2 = nequal t1 t2
945 equal :: Eq a => IntMap a -> IntMap a -> Bool
946 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
947 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
948 equal (Tip kx x) (Tip ky y)
949 = (kx == ky) && (x==y)
953 nequal :: Eq a => IntMap a -> IntMap a -> Bool
954 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
955 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
956 nequal (Tip kx x) (Tip ky y)
957 = (kx /= ky) || (x/=y)
958 nequal Nil Nil = False
961 {--------------------------------------------------------------------
963 --------------------------------------------------------------------}
965 instance Ord a => Ord (IntMap a) where
966 compare m1 m2 = compare (toList m1) (toList m2)
968 {--------------------------------------------------------------------
970 --------------------------------------------------------------------}
972 instance Functor IntMap where
975 {--------------------------------------------------------------------
977 --------------------------------------------------------------------}
979 instance Ord a => Monoid (IntMap a) where
984 {--------------------------------------------------------------------
986 --------------------------------------------------------------------}
988 instance Show a => Show (IntMap a) where
989 showsPrec d t = showMap (toList t)
992 showMap :: (Show a) => [(Key,a)] -> ShowS
996 = showChar '{' . showElem x . showTail xs
998 showTail [] = showChar '}'
999 showTail (x:xs) = showChar ',' . showElem x . showTail xs
1001 showElem (k,x) = shows k . showString ":=" . shows x
1003 {--------------------------------------------------------------------
1005 --------------------------------------------------------------------}
1006 -- | /O(n)/. Show the tree that implements the map. The tree is shown
1007 -- in a compressed, hanging format.
1008 showTree :: Show a => IntMap a -> String
1010 = showTreeWith True False s
1013 {- | /O(n)/. The expression (@showTreeWith hang wide map@) shows
1014 the tree that implements the map. If @hang@ is
1015 @True@, a /hanging/ tree is shown otherwise a rotated tree is shown. If
1016 @wide@ is true, an extra wide version is shown.
1018 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
1019 showTreeWith hang wide t
1020 | hang = (showsTreeHang wide [] t) ""
1021 | otherwise = (showsTree wide [] [] t) ""
1023 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
1024 showsTree wide lbars rbars t
1027 -> showsTree wide (withBar rbars) (withEmpty rbars) r .
1028 showWide wide rbars .
1029 showsBars lbars . showString (showBin p m) . showString "\n" .
1030 showWide wide lbars .
1031 showsTree wide (withEmpty lbars) (withBar lbars) l
1033 -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1034 Nil -> showsBars lbars . showString "|\n"
1036 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
1037 showsTreeHang wide bars t
1040 -> showsBars bars . showString (showBin p m) . showString "\n" .
1041 showWide wide bars .
1042 showsTreeHang wide (withBar bars) l .
1043 showWide wide bars .
1044 showsTreeHang wide (withEmpty bars) r
1046 -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1047 Nil -> showsBars bars . showString "|\n"
1050 = "*" -- ++ show (p,m)
1053 | wide = showString (concat (reverse bars)) . showString "|\n"
1056 showsBars :: [String] -> ShowS
1060 _ -> showString (concat (reverse (tail bars))) . showString node
1063 withBar bars = "| ":bars
1064 withEmpty bars = " ":bars
1067 {--------------------------------------------------------------------
1069 --------------------------------------------------------------------}
1070 {--------------------------------------------------------------------
1072 --------------------------------------------------------------------}
1073 join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
1075 | zero p1 m = Bin p m t1 t2
1076 | otherwise = Bin p m t2 t1
1078 m = branchMask p1 p2
1081 {--------------------------------------------------------------------
1082 @bin@ assures that we never have empty trees within a tree.
1083 --------------------------------------------------------------------}
1084 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
1087 bin p m l r = Bin p m l r
1090 {--------------------------------------------------------------------
1091 Endian independent bit twiddling
1092 --------------------------------------------------------------------}
1093 zero :: Key -> Mask -> Bool
1095 = (natFromInt i) .&. (natFromInt m) == 0
1097 nomatch,match :: Key -> Prefix -> Mask -> Bool
1104 mask :: Key -> Mask -> Prefix
1106 = maskW (natFromInt i) (natFromInt m)
1109 zeroN :: Nat -> Nat -> Bool
1110 zeroN i m = (i .&. m) == 0
1112 {--------------------------------------------------------------------
1113 Big endian operations
1114 --------------------------------------------------------------------}
1115 maskW :: Nat -> Nat -> Prefix
1117 = intFromNat (i .&. (complement (m-1) `xor` m))
1119 shorter :: Mask -> Mask -> Bool
1121 = (natFromInt m1) > (natFromInt m2)
1123 branchMask :: Prefix -> Prefix -> Mask
1125 = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
1127 {----------------------------------------------------------------------
1128 Finding the highest bit (mask) in a word [x] can be done efficiently in
1130 * convert to a floating point value and the mantissa tells us the
1131 [log2(x)] that corresponds with the highest bit position. The mantissa
1132 is retrieved either via the standard C function [frexp] or by some bit
1133 twiddling on IEEE compatible numbers (float). Note that one needs to
1134 use at least [double] precision for an accurate mantissa of 32 bit
1136 * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
1137 * use processor specific assembler instruction (asm).
1139 The most portable way would be [bit], but is it efficient enough?
1140 I have measured the cycle counts of the different methods on an AMD
1141 Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
1143 highestBitMask: method cycles
1150 highestBit: method cycles
1157 Wow, the bit twiddling is on today's RISC like machines even faster
1158 than a single CISC instruction (BSR)!
1159 ----------------------------------------------------------------------}
1161 {----------------------------------------------------------------------
1162 [highestBitMask] returns a word where only the highest bit is set.
1163 It is found by first setting all bits in lower positions than the
1164 highest bit and than taking an exclusive or with the original value.
1165 Allthough the function may look expensive, GHC compiles this into
1166 excellent C code that subsequently compiled into highly efficient
1167 machine code. The algorithm is derived from Jorg Arndt's FXT library.
1168 ----------------------------------------------------------------------}
1169 highestBitMask :: Nat -> Nat
1171 = case (x .|. shiftRL x 1) of
1172 x -> case (x .|. shiftRL x 2) of
1173 x -> case (x .|. shiftRL x 4) of
1174 x -> case (x .|. shiftRL x 8) of
1175 x -> case (x .|. shiftRL x 16) of
1176 x -> case (x .|. shiftRL x 32) of -- for 64 bit platforms
1177 x -> (x `xor` (shiftRL x 1))
1180 {--------------------------------------------------------------------
1182 --------------------------------------------------------------------}
1186 (x:xx) -> let z' = f z x in seq z' (foldlStrict f z' xx)
1189 {--------------------------------------------------------------------
1191 --------------------------------------------------------------------}
1192 testTree :: [Int] -> IntMap Int
1193 testTree xs = fromList [(x,x*x*30696 `mod` 65521) | x <- xs]
1194 test1 = testTree [1..20]
1195 test2 = testTree [30,29..10]
1196 test3 = testTree [1,4,6,89,2323,53,43,234,5,79,12,9,24,9,8,423,8,42,4,8,9,3]
1198 {--------------------------------------------------------------------
1200 --------------------------------------------------------------------}
1205 { configMaxTest = 500
1206 , configMaxFail = 5000
1207 , configSize = \n -> (div n 2 + 3)
1208 , configEvery = \n args -> let s = show n in s ++ [ '\b' | _ <- s ]
1212 {--------------------------------------------------------------------
1213 Arbitrary, reasonably balanced trees
1214 --------------------------------------------------------------------}
1215 instance Arbitrary a => Arbitrary (IntMap a) where
1216 arbitrary = do{ ks <- arbitrary
1217 ; xs <- mapM (\k -> do{ x <- arbitrary; return (k,x)}) ks
1218 ; return (fromList xs)
1222 {--------------------------------------------------------------------
1223 Single, Insert, Delete
1224 --------------------------------------------------------------------}
1225 prop_Single :: Key -> Int -> Bool
1227 = (insert k x empty == singleton k x)
1229 prop_InsertDelete :: Key -> Int -> IntMap Int -> Property
1230 prop_InsertDelete k x t
1231 = not (member k t) ==> delete k (insert k x t) == t
1233 prop_UpdateDelete :: Key -> IntMap Int -> Bool
1234 prop_UpdateDelete k t
1235 = update (const Nothing) k t == delete k t
1238 {--------------------------------------------------------------------
1240 --------------------------------------------------------------------}
1241 prop_UnionInsert :: Key -> Int -> IntMap Int -> Bool
1242 prop_UnionInsert k x t
1243 = union (singleton k x) t == insert k x t
1245 prop_UnionAssoc :: IntMap Int -> IntMap Int -> IntMap Int -> Bool
1246 prop_UnionAssoc t1 t2 t3
1247 = union t1 (union t2 t3) == union (union t1 t2) t3
1249 prop_UnionComm :: IntMap Int -> IntMap Int -> Bool
1250 prop_UnionComm t1 t2
1251 = (union t1 t2 == unionWith (\x y -> y) t2 t1)
1254 prop_Diff :: [(Key,Int)] -> [(Key,Int)] -> Bool
1256 = List.sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys)))
1257 == List.sort ((List.\\) (nub (Prelude.map fst xs)) (nub (Prelude.map fst ys)))
1259 prop_Int :: [(Key,Int)] -> [(Key,Int)] -> Bool
1261 = List.sort (keys (intersection (fromListWith (+) xs) (fromListWith (+) ys)))
1262 == List.sort (nub ((List.intersect) (Prelude.map fst xs) (Prelude.map fst ys)))
1264 {--------------------------------------------------------------------
1266 --------------------------------------------------------------------}
1268 = forAll (choose (5,100)) $ \n ->
1269 let xs = [(x,()) | x <- [0..n::Int]]
1270 in fromAscList xs == fromList xs
1272 prop_List :: [Key] -> Bool
1274 = (sort (nub xs) == [x | (x,()) <- toAscList (fromList [(x,()) | x <- xs])])