1 {-# OPTIONS -cpp -fglasgow-exts #-}
2 -----------------------------------------------------------------------------
3 -- Module : Data.IntMap
4 -- Copyright : (c) Daan Leijen 2002
6 -- Maintainer : libraries@haskell.org
7 -- Stability : provisional
8 -- Portability : portable
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
18 -- structure performs especially well on binary operations like 'union'
19 -- and 'intersection'. However, my benchmarks show that it is also
20 -- (much) faster on insertions and deletions when compared to a generic
21 -- size-balanced map implementation (see "Data.Map" and "Data.FiniteMap").
23 -- * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\",
24 -- Workshop on ML, September 1998, pages 77-86,
25 -- <http://www.cse.ogi.edu/~andy/pub/finite.htm>
27 -- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve
28 -- Information Coded In Alphanumeric/\", Journal of the ACM, 15(4),
29 -- October 1968, pages 514-534.
31 -- Many operations have a worst-case complexity of /O(min(n,W))/.
32 -- This means that the operation can become linear in the number of
33 -- elements with a maximum of /W/ -- the number of bits in an 'Int'
35 -----------------------------------------------------------------------------
39 IntMap, Key -- instance Eq,Show
57 , insertWith, insertWithKey, insertLookupWithKey
114 , fromDistinctAscList
126 , isSubmapOf, isSubmapOfBy
127 , isProperSubmapOf, isProperSubmapOfBy
135 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
139 import qualified Data.IntSet as IntSet
143 import qualified Prelude
144 import Debug.QuickCheck
145 import List (nub,sort)
146 import qualified List
149 #if __GLASGOW_HASKELL__ >= 503
151 import GHC.Exts ( Word(..), Int(..), shiftRL# )
152 #elif __GLASGOW_HASKELL__
154 import GlaExts ( Word(..), Int(..), shiftRL# )
159 infixl 9 \\{-This comment teaches CPP correct behaviour -}
161 -- A "Nat" is a natural machine word (an unsigned Int)
164 natFromInt :: Key -> Nat
165 natFromInt i = fromIntegral i
167 intFromNat :: Nat -> Key
168 intFromNat w = fromIntegral w
170 shiftRL :: Nat -> Key -> Nat
171 #if __GLASGOW_HASKELL__
172 {--------------------------------------------------------------------
173 GHC: use unboxing to get @shiftRL@ inlined.
174 --------------------------------------------------------------------}
175 shiftRL (W# x) (I# i)
178 shiftRL x i = shiftR x i
181 {--------------------------------------------------------------------
183 --------------------------------------------------------------------}
185 -- | /O(min(n,W))/. Find the value of a key. Calls @error@ when the element can not be found.
187 (!) :: IntMap a -> Key -> a
190 -- | /O(n+m)/. See 'difference'.
191 (\\) :: IntMap a -> IntMap b -> IntMap a
192 m1 \\ m2 = difference m1 m2
194 {--------------------------------------------------------------------
196 --------------------------------------------------------------------}
197 -- | A map of integers to values @a@.
199 | Tip {-# UNPACK #-} !Key a
200 | Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(IntMap a) !(IntMap a)
206 {--------------------------------------------------------------------
208 --------------------------------------------------------------------}
209 -- | /O(1)/. Is the map empty?
210 null :: IntMap a -> Bool
214 -- | /O(n)/. Number of elements in the map.
215 size :: IntMap a -> Int
218 Bin p m l r -> size l + size r
222 -- | /O(min(n,W))/. Is the key a member of the map?
223 member :: Key -> IntMap a -> Bool
229 -- | /O(min(n,W))/. Lookup the value of a key in the map.
230 lookup :: Key -> IntMap a -> Maybe a
232 = let nk = natFromInt k in seq nk (lookupN nk t)
234 lookupN :: Nat -> IntMap a -> Maybe a
238 | zeroN k (natFromInt m) -> lookupN k l
239 | otherwise -> lookupN k r
241 | (k == natFromInt kx) -> Just x
242 | otherwise -> Nothing
245 find' :: Key -> IntMap a -> a
248 Nothing -> error ("IntMap.find: key " ++ show k ++ " is not an element of the map")
252 -- | /O(min(n,W))/. The expression @(findWithDefault def k map)@ returns the value of key @k@ or returns @def@ when
253 -- the key is not an element of the map.
254 findWithDefault :: a -> Key -> IntMap a -> a
255 findWithDefault def k m
260 {--------------------------------------------------------------------
262 --------------------------------------------------------------------}
263 -- | /O(1)/. The empty map.
268 -- | /O(1)/. A map of one element.
269 singleton :: Key -> a -> IntMap a
273 {--------------------------------------------------------------------
275 'insert' is the inlined version of 'insertWith (\k x y -> x)'
276 --------------------------------------------------------------------}
277 -- | /O(min(n,W))/. Insert a new key\/value pair in the map. When the key
278 -- is already an element of the set, its value is replaced by the new value,
279 -- ie. 'insert' is left-biased.
280 insert :: Key -> a -> IntMap a -> IntMap a
284 | nomatch k p m -> join k (Tip k x) p t
285 | zero k m -> Bin p m (insert k x l) r
286 | otherwise -> Bin p m l (insert k x r)
289 | otherwise -> join k (Tip k x) ky t
292 -- right-biased insertion, used by 'union'
293 -- | /O(min(n,W))/. Insert with a combining function.
294 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
296 = insertWithKey (\k x y -> f x y) k x t
298 -- | /O(min(n,W))/. Insert with a combining function.
299 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
300 insertWithKey f k x t
303 | nomatch k p m -> join k (Tip k x) p t
304 | zero k m -> Bin p m (insertWithKey f k x l) r
305 | otherwise -> Bin p m l (insertWithKey f k x r)
307 | k==ky -> Tip k (f k x y)
308 | otherwise -> join k (Tip k x) ky t
312 -- | /O(min(n,W))/. The expression (@insertLookupWithKey f k x map@) is a pair where
313 -- the first element is equal to (@lookup k map@) and the second element
314 -- equal to (@insertWithKey f k x map@).
315 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
316 insertLookupWithKey f k x t
319 | nomatch k p m -> (Nothing,join k (Tip k x) p t)
320 | zero k m -> let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r)
321 | otherwise -> let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r')
323 | k==ky -> (Just y,Tip k (f k x y))
324 | otherwise -> (Nothing,join k (Tip k x) ky t)
325 Nil -> (Nothing,Tip k x)
328 {--------------------------------------------------------------------
330 [delete] is the inlined version of [deleteWith (\k x -> Nothing)]
331 --------------------------------------------------------------------}
332 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
333 -- a member of the map, the original map is returned.
334 delete :: Key -> IntMap a -> IntMap a
339 | zero k m -> bin p m (delete k l) r
340 | otherwise -> bin p m l (delete k r)
346 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
347 -- a member of the map, the original map is returned.
348 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
350 = adjustWithKey (\k x -> f x) k m
352 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
353 -- a member of the map, the original map is returned.
354 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
356 = updateWithKey (\k x -> Just (f k x)) k m
358 -- | /O(min(n,W))/. The expression (@update f k map@) updates the value @x@
359 -- at @k@ (if it is in the map). If (@f x@) is @Nothing@, the element is
360 -- deleted. If it is (@Just y@), the key @k@ is bound to the new value @y@.
361 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
363 = updateWithKey (\k x -> f x) k m
365 -- | /O(min(n,W))/. The expression (@update f k map@) updates the value @x@
366 -- at @k@ (if it is in the map). If (@f k x@) is @Nothing@, the element is
367 -- deleted. If it is (@Just y@), the key @k@ is bound to the new value @y@.
368 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
373 | zero k m -> bin p m (updateWithKey f k l) r
374 | otherwise -> bin p m l (updateWithKey f k r)
376 | k==ky -> case (f k y) of
382 -- | /O(min(n,W))/. Lookup and update.
383 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
384 updateLookupWithKey f k t
387 | nomatch k p m -> (Nothing,t)
388 | zero k m -> let (found,l') = updateLookupWithKey f k l in (found,bin p m l' r)
389 | otherwise -> let (found,r') = updateLookupWithKey f k r in (found,bin p m l r')
391 | k==ky -> case (f k y) of
392 Just y' -> (Just y,Tip ky y')
393 Nothing -> (Just y,Nil)
394 | otherwise -> (Nothing,t)
398 {--------------------------------------------------------------------
400 --------------------------------------------------------------------}
401 -- | The union of a list of maps.
402 unions :: [IntMap a] -> IntMap a
404 = foldlStrict union empty xs
406 -- | The union of a list of maps, with a combining operation
407 unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a
409 = foldlStrict (unionWith f) empty ts
411 -- | /O(n+m)/. The (left-biased) union of two sets.
412 union :: IntMap a -> IntMap a -> IntMap a
413 union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
414 | shorter m1 m2 = union1
415 | shorter m2 m1 = union2
416 | p1 == p2 = Bin p1 m1 (union l1 l2) (union r1 r2)
417 | otherwise = join p1 t1 p2 t2
419 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
420 | zero p2 m1 = Bin p1 m1 (union l1 t2) r1
421 | otherwise = Bin p1 m1 l1 (union r1 t2)
423 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
424 | zero p1 m2 = Bin p2 m2 (union t1 l2) r2
425 | otherwise = Bin p2 m2 l2 (union t1 r2)
427 union (Tip k x) t = insert k x t
428 union t (Tip k x) = insertWith (\x y -> y) k x t -- right bias
432 -- | /O(n+m)/. The union with a combining function.
433 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
435 = unionWithKey (\k x y -> f x y) m1 m2
437 -- | /O(n+m)/. The union with a combining function.
438 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
439 unionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
440 | shorter m1 m2 = union1
441 | shorter m2 m1 = union2
442 | p1 == p2 = Bin p1 m1 (unionWithKey f l1 l2) (unionWithKey f r1 r2)
443 | otherwise = join p1 t1 p2 t2
445 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
446 | zero p2 m1 = Bin p1 m1 (unionWithKey f l1 t2) r1
447 | otherwise = Bin p1 m1 l1 (unionWithKey f r1 t2)
449 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
450 | zero p1 m2 = Bin p2 m2 (unionWithKey f t1 l2) r2
451 | otherwise = Bin p2 m2 l2 (unionWithKey f t1 r2)
453 unionWithKey f (Tip k x) t = insertWithKey f k x t
454 unionWithKey f t (Tip k x) = insertWithKey (\k x y -> f k y x) k x t -- right bias
455 unionWithKey f Nil t = t
456 unionWithKey f t Nil = t
458 {--------------------------------------------------------------------
460 --------------------------------------------------------------------}
461 -- | /O(n+m)/. Difference between two maps (based on keys).
462 difference :: IntMap a -> IntMap b -> IntMap a
463 difference t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
464 | shorter m1 m2 = difference1
465 | shorter m2 m1 = difference2
466 | p1 == p2 = bin p1 m1 (difference l1 l2) (difference r1 r2)
469 difference1 | nomatch p2 p1 m1 = t1
470 | zero p2 m1 = bin p1 m1 (difference l1 t2) r1
471 | otherwise = bin p1 m1 l1 (difference r1 t2)
473 difference2 | nomatch p1 p2 m2 = t1
474 | zero p1 m2 = difference t1 l2
475 | otherwise = difference t1 r2
477 difference t1@(Tip k x) t2
481 difference Nil t = Nil
482 difference t (Tip k x) = delete k t
485 -- | /O(n+m)/. Difference with a combining function.
486 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
487 differenceWith f m1 m2
488 = differenceWithKey (\k x y -> f x y) m1 m2
490 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
491 -- encountered, the combining function is applied to the key and both values.
492 -- If it returns @Nothing@, the element is discarded (proper set difference). If
493 -- it returns (@Just y@), the element is updated with a new value @y@.
494 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
495 differenceWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
496 | shorter m1 m2 = difference1
497 | shorter m2 m1 = difference2
498 | p1 == p2 = bin p1 m1 (differenceWithKey f l1 l2) (differenceWithKey f r1 r2)
501 difference1 | nomatch p2 p1 m1 = t1
502 | zero p2 m1 = bin p1 m1 (differenceWithKey f l1 t2) r1
503 | otherwise = bin p1 m1 l1 (differenceWithKey f r1 t2)
505 difference2 | nomatch p1 p2 m2 = t1
506 | zero p1 m2 = differenceWithKey f t1 l2
507 | otherwise = differenceWithKey f t1 r2
509 differenceWithKey f t1@(Tip k x) t2
510 = case lookup k t2 of
511 Just y -> case f k x y of
516 differenceWithKey f Nil t = Nil
517 differenceWithKey f t (Tip k y) = updateWithKey (\k x -> f k x y) k t
518 differenceWithKey f t Nil = t
521 {--------------------------------------------------------------------
523 --------------------------------------------------------------------}
524 -- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys).
525 intersection :: IntMap a -> IntMap b -> IntMap a
526 intersection t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
527 | shorter m1 m2 = intersection1
528 | shorter m2 m1 = intersection2
529 | p1 == p2 = bin p1 m1 (intersection l1 l2) (intersection r1 r2)
532 intersection1 | nomatch p2 p1 m1 = Nil
533 | zero p2 m1 = intersection l1 t2
534 | otherwise = intersection r1 t2
536 intersection2 | nomatch p1 p2 m2 = Nil
537 | zero p1 m2 = intersection t1 l2
538 | otherwise = intersection t1 r2
540 intersection t1@(Tip k x) t2
543 intersection t (Tip k x)
547 intersection Nil t = Nil
548 intersection t Nil = Nil
550 -- | /O(n+m)/. The intersection with a combining function.
551 intersectionWith :: (a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
552 intersectionWith f m1 m2
553 = intersectionWithKey (\k x y -> f x y) m1 m2
555 -- | /O(n+m)/. The intersection with a combining function.
556 intersectionWithKey :: (Key -> a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
557 intersectionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
558 | shorter m1 m2 = intersection1
559 | shorter m2 m1 = intersection2
560 | p1 == p2 = bin p1 m1 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2)
563 intersection1 | nomatch p2 p1 m1 = Nil
564 | zero p2 m1 = intersectionWithKey f l1 t2
565 | otherwise = intersectionWithKey f r1 t2
567 intersection2 | nomatch p1 p2 m2 = Nil
568 | zero p1 m2 = intersectionWithKey f t1 l2
569 | otherwise = intersectionWithKey f t1 r2
571 intersectionWithKey f t1@(Tip k x) t2
572 = case lookup k t2 of
573 Just y -> Tip k (f k x y)
575 intersectionWithKey f t1 (Tip k y)
576 = case lookup k t1 of
577 Just x -> Tip k (f k x y)
579 intersectionWithKey f Nil t = Nil
580 intersectionWithKey f t Nil = Nil
583 {--------------------------------------------------------------------
585 --------------------------------------------------------------------}
586 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
587 -- Defined as (@isProperSubmapOf = isProperSubmapOfBy (==)@).
588 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
589 isProperSubmapOf m1 m2
590 = isProperSubmapOfBy (==) m1 m2
592 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
593 The expression (@isProperSubmapOfBy f m1 m2@) returns @True@ when
594 @m1@ and @m2@ are not equal,
595 all keys in @m1@ are in @m2@, and when @f@ returns @True@ when
596 applied to their respective values. For example, the following
597 expressions are all @True@.
599 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
600 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
602 But the following are all @False@:
604 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
605 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
606 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
608 isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
609 isProperSubmapOfBy pred t1 t2
610 = case submapCmp pred t1 t2 of
614 submapCmp pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
616 | shorter m2 m1 = submapCmpLt
617 | p1 == p2 = submapCmpEq
618 | otherwise = GT -- disjoint
620 submapCmpLt | nomatch p1 p2 m2 = GT
621 | zero p1 m2 = submapCmp pred t1 l2
622 | otherwise = submapCmp pred t1 r2
623 submapCmpEq = case (submapCmp pred l1 l2, submapCmp pred r1 r2) of
629 submapCmp pred (Bin p m l r) t = GT
630 submapCmp pred (Tip kx x) (Tip ky y)
631 | (kx == ky) && pred x y = EQ
632 | otherwise = GT -- disjoint
633 submapCmp pred (Tip k x) t
635 Just y | pred x y -> LT
636 other -> GT -- disjoint
637 submapCmp pred Nil Nil = EQ
638 submapCmp pred Nil t = LT
640 -- | /O(n+m)/. Is this a submap? Defined as (@isSubmapOf = isSubmapOfBy (==)@).
641 isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
643 = isSubmapOfBy (==) m1 m2
646 The expression (@isSubmapOfBy f m1 m2@) returns @True@ if
647 all keys in @m1@ are in @m2@, and when @f@ returns @True@ when
648 applied to their respective values. For example, the following
649 expressions are all @True@.
651 > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
652 > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
653 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
655 But the following are all @False@:
657 > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)])
658 > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
659 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
662 isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
663 isSubmapOfBy pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
664 | shorter m1 m2 = False
665 | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy pred t1 l2
666 else isSubmapOfBy pred t1 r2)
667 | otherwise = (p1==p2) && isSubmapOfBy pred l1 l2 && isSubmapOfBy pred r1 r2
668 isSubmapOfBy pred (Bin p m l r) t = False
669 isSubmapOfBy pred (Tip k x) t = case lookup k t of
672 isSubmapOfBy pred Nil t = True
674 {--------------------------------------------------------------------
676 --------------------------------------------------------------------}
677 -- | /O(n)/. Map a function over all values in the map.
678 map :: (a -> b) -> IntMap a -> IntMap b
680 = mapWithKey (\k x -> f x) m
682 -- | /O(n)/. Map a function over all values in the map.
683 mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b
686 Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r)
687 Tip k x -> Tip k (f k x)
690 -- | /O(n)/. The function @mapAccum@ threads an accumulating
691 -- argument through the map in an unspecified order.
692 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
694 = mapAccumWithKey (\a k x -> f a x) a m
696 -- | /O(n)/. The function @mapAccumWithKey@ threads an accumulating
697 -- argument through the map in an unspecified order.
698 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
699 mapAccumWithKey f a t
702 -- | /O(n)/. The function @mapAccumL@ threads an accumulating
703 -- argument through the map in pre-order.
704 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
707 Bin p m l r -> let (a1,l') = mapAccumL f a l
708 (a2,r') = mapAccumL f a1 r
709 in (a2,Bin p m l' r')
710 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
714 -- | /O(n)/. The function @mapAccumR@ threads an accumulating
715 -- argument throught the map in post-order.
716 mapAccumR :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
719 Bin p m l r -> let (a1,r') = mapAccumR f a r
720 (a2,l') = mapAccumR f a1 l
721 in (a2,Bin p m l' r')
722 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
725 {--------------------------------------------------------------------
727 --------------------------------------------------------------------}
728 -- | /O(n)/. Filter all values that satisfy some predicate.
729 filter :: (a -> Bool) -> IntMap a -> IntMap a
731 = filterWithKey (\k x -> p x) m
733 -- | /O(n)/. Filter all keys\/values that satisfy some predicate.
734 filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a
738 -> bin p m (filterWithKey pred l) (filterWithKey pred r)
744 -- | /O(n)/. partition the map according to some predicate. The first
745 -- map contains all elements that satisfy the predicate, the second all
746 -- elements that fail the predicate. See also 'split'.
747 partition :: (a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
749 = partitionWithKey (\k x -> p x) m
751 -- | /O(n)/. partition the map according to some predicate. The first
752 -- map contains all elements that satisfy the predicate, the second all
753 -- elements that fail the predicate. See also 'split'.
754 partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
755 partitionWithKey pred t
758 -> let (l1,l2) = partitionWithKey pred l
759 (r1,r2) = partitionWithKey pred r
760 in (bin p m l1 r1, bin p m l2 r2)
762 | pred k x -> (t,Nil)
763 | otherwise -> (Nil,t)
767 -- | /O(log n)/. The expression (@split k map@) is a pair @(map1,map2)@
768 -- where all keys in @map1@ are lower than @k@ and all keys in
769 -- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
770 split :: Key -> IntMap a -> (IntMap a,IntMap a)
774 | zero k m -> let (lt,gt) = split k l in (lt,union gt r)
775 | otherwise -> let (lt,gt) = split k r in (union l lt,gt)
779 | otherwise -> (Nil,Nil)
782 -- | /O(log n)/. Performs a 'split' but also returns whether the pivot
783 -- key was found in the original map.
784 splitLookup :: Key -> IntMap a -> (Maybe a,IntMap a,IntMap a)
788 | zero k m -> let (found,lt,gt) = splitLookup k l in (found,lt,union gt r)
789 | otherwise -> let (found,lt,gt) = splitLookup k r in (found,union l lt,gt)
791 | k>ky -> (Nothing,t,Nil)
792 | k<ky -> (Nothing,Nil,t)
793 | otherwise -> (Just y,Nil,Nil)
794 Nil -> (Nothing,Nil,Nil)
796 {--------------------------------------------------------------------
798 --------------------------------------------------------------------}
799 -- | /O(n)/. Fold over the elements of a map in an unspecified order.
801 -- > sum map = fold (+) 0 map
802 -- > elems map = fold (:) [] map
803 fold :: (a -> b -> b) -> b -> IntMap a -> b
805 = foldWithKey (\k x y -> f x y) z t
807 -- | /O(n)/. Fold over the elements of a map in an unspecified order.
809 -- > keys map = foldWithKey (\k x ks -> k:ks) [] map
810 foldWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
814 foldr :: (Key -> a -> b -> b) -> b -> IntMap a -> b
817 Bin p m l r -> foldr f (foldr f z r) l
821 {--------------------------------------------------------------------
823 --------------------------------------------------------------------}
824 -- | /O(n)/. Return all elements of the map.
825 elems :: IntMap a -> [a]
827 = foldWithKey (\k x xs -> x:xs) [] m
829 -- | /O(n)/. Return all keys of the map.
830 keys :: IntMap a -> [Key]
832 = foldWithKey (\k x ks -> k:ks) [] m
834 -- | /O(n*min(n,W))/. The set of all keys of the map.
835 keysSet :: IntMap a -> IntSet.IntSet
836 keysSet m = IntSet.fromDistinctAscList (keys m)
839 -- | /O(n)/. Return all key\/value pairs in the map.
840 assocs :: IntMap a -> [(Key,a)]
845 {--------------------------------------------------------------------
847 --------------------------------------------------------------------}
848 -- | /O(n)/. Convert the map to a list of key\/value pairs.
849 toList :: IntMap a -> [(Key,a)]
851 = foldWithKey (\k x xs -> (k,x):xs) [] t
853 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
854 -- keys are in ascending order.
855 toAscList :: IntMap a -> [(Key,a)]
857 = -- NOTE: the following algorithm only works for big-endian trees
858 let (pos,neg) = span (\(k,x) -> k >=0) (foldr (\k x xs -> (k,x):xs) [] t) in neg ++ pos
860 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
861 fromList :: [(Key,a)] -> IntMap a
863 = foldlStrict ins empty xs
865 ins t (k,x) = insert k x t
867 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
868 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
870 = fromListWithKey (\k x y -> f x y) xs
872 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
873 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
875 = foldlStrict ins empty xs
877 ins t (k,x) = insertWithKey f k x t
879 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
880 -- the keys are in ascending order.
881 fromAscList :: [(Key,a)] -> IntMap a
885 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
886 -- the keys are in ascending order, with a combining function on equal keys.
887 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
891 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
892 -- the keys are in ascending order, with a combining function on equal keys.
893 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
894 fromAscListWithKey f xs
895 = fromListWithKey f xs
897 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
898 -- the keys are in ascending order and all distinct.
899 fromDistinctAscList :: [(Key,a)] -> IntMap a
900 fromDistinctAscList xs
904 {--------------------------------------------------------------------
906 --------------------------------------------------------------------}
907 instance Eq a => Eq (IntMap a) where
908 t1 == t2 = equal t1 t2
909 t1 /= t2 = nequal t1 t2
911 equal :: Eq a => IntMap a -> IntMap a -> Bool
912 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
913 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
914 equal (Tip kx x) (Tip ky y)
915 = (kx == ky) && (x==y)
919 nequal :: Eq a => IntMap a -> IntMap a -> Bool
920 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
921 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
922 nequal (Tip kx x) (Tip ky y)
923 = (kx /= ky) || (x/=y)
924 nequal Nil Nil = False
927 {--------------------------------------------------------------------
929 --------------------------------------------------------------------}
931 instance Ord a => Ord (IntMap a) where
932 compare m1 m2 = compare (toList m1) (toList m2)
934 {--------------------------------------------------------------------
936 --------------------------------------------------------------------}
938 instance Functor IntMap where
941 {--------------------------------------------------------------------
943 --------------------------------------------------------------------}
945 instance Ord a => Monoid (IntMap a) where
950 {--------------------------------------------------------------------
952 --------------------------------------------------------------------}
954 instance Show a => Show (IntMap a) where
955 showsPrec d t = showMap (toList t)
958 showMap :: (Show a) => [(Key,a)] -> ShowS
962 = showChar '{' . showElem x . showTail xs
964 showTail [] = showChar '}'
965 showTail (x:xs) = showChar ',' . showElem x . showTail xs
967 showElem (k,x) = shows k . showString ":=" . shows x
969 {--------------------------------------------------------------------
971 --------------------------------------------------------------------}
972 -- | /O(n)/. Show the tree that implements the map. The tree is shown
973 -- in a compressed, hanging format.
974 showTree :: Show a => IntMap a -> String
976 = showTreeWith True False s
979 {- | /O(n)/. The expression (@showTreeWith hang wide map@) shows
980 the tree that implements the map. If @hang@ is
981 @True@, a /hanging/ tree is shown otherwise a rotated tree is shown. If
982 @wide@ is true, an extra wide version is shown.
984 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
985 showTreeWith hang wide t
986 | hang = (showsTreeHang wide [] t) ""
987 | otherwise = (showsTree wide [] [] t) ""
989 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
990 showsTree wide lbars rbars t
993 -> showsTree wide (withBar rbars) (withEmpty rbars) r .
994 showWide wide rbars .
995 showsBars lbars . showString (showBin p m) . showString "\n" .
996 showWide wide lbars .
997 showsTree wide (withEmpty lbars) (withBar lbars) l
999 -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1000 Nil -> showsBars lbars . showString "|\n"
1002 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
1003 showsTreeHang wide bars t
1006 -> showsBars bars . showString (showBin p m) . showString "\n" .
1007 showWide wide bars .
1008 showsTreeHang wide (withBar bars) l .
1009 showWide wide bars .
1010 showsTreeHang wide (withEmpty bars) r
1012 -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1013 Nil -> showsBars bars . showString "|\n"
1016 = "*" -- ++ show (p,m)
1019 | wide = showString (concat (reverse bars)) . showString "|\n"
1022 showsBars :: [String] -> ShowS
1026 _ -> showString (concat (reverse (tail bars))) . showString node
1029 withBar bars = "| ":bars
1030 withEmpty bars = " ":bars
1033 {--------------------------------------------------------------------
1035 --------------------------------------------------------------------}
1036 {--------------------------------------------------------------------
1038 --------------------------------------------------------------------}
1039 join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
1041 | zero p1 m = Bin p m t1 t2
1042 | otherwise = Bin p m t2 t1
1044 m = branchMask p1 p2
1047 {--------------------------------------------------------------------
1048 @bin@ assures that we never have empty trees within a tree.
1049 --------------------------------------------------------------------}
1050 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
1053 bin p m l r = Bin p m l r
1056 {--------------------------------------------------------------------
1057 Endian independent bit twiddling
1058 --------------------------------------------------------------------}
1059 zero :: Key -> Mask -> Bool
1061 = (natFromInt i) .&. (natFromInt m) == 0
1063 nomatch,match :: Key -> Prefix -> Mask -> Bool
1070 mask :: Key -> Mask -> Prefix
1072 = maskW (natFromInt i) (natFromInt m)
1075 zeroN :: Nat -> Nat -> Bool
1076 zeroN i m = (i .&. m) == 0
1078 {--------------------------------------------------------------------
1079 Big endian operations
1080 --------------------------------------------------------------------}
1081 maskW :: Nat -> Nat -> Prefix
1083 = intFromNat (i .&. (complement (m-1) `xor` m))
1085 shorter :: Mask -> Mask -> Bool
1087 = (natFromInt m1) > (natFromInt m2)
1089 branchMask :: Prefix -> Prefix -> Mask
1091 = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
1093 {----------------------------------------------------------------------
1094 Finding the highest bit (mask) in a word [x] can be done efficiently in
1096 * convert to a floating point value and the mantissa tells us the
1097 [log2(x)] that corresponds with the highest bit position. The mantissa
1098 is retrieved either via the standard C function [frexp] or by some bit
1099 twiddling on IEEE compatible numbers (float). Note that one needs to
1100 use at least [double] precision for an accurate mantissa of 32 bit
1102 * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
1103 * use processor specific assembler instruction (asm).
1105 The most portable way would be [bit], but is it efficient enough?
1106 I have measured the cycle counts of the different methods on an AMD
1107 Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
1109 highestBitMask: method cycles
1116 highestBit: method cycles
1123 Wow, the bit twiddling is on today's RISC like machines even faster
1124 than a single CISC instruction (BSR)!
1125 ----------------------------------------------------------------------}
1127 {----------------------------------------------------------------------
1128 [highestBitMask] returns a word where only the highest bit is set.
1129 It is found by first setting all bits in lower positions than the
1130 highest bit and than taking an exclusive or with the original value.
1131 Allthough the function may look expensive, GHC compiles this into
1132 excellent C code that subsequently compiled into highly efficient
1133 machine code. The algorithm is derived from Jorg Arndt's FXT library.
1134 ----------------------------------------------------------------------}
1135 highestBitMask :: Nat -> Nat
1137 = case (x .|. shiftRL x 1) of
1138 x -> case (x .|. shiftRL x 2) of
1139 x -> case (x .|. shiftRL x 4) of
1140 x -> case (x .|. shiftRL x 8) of
1141 x -> case (x .|. shiftRL x 16) of
1142 x -> case (x .|. shiftRL x 32) of -- for 64 bit platforms
1143 x -> (x `xor` (shiftRL x 1))
1146 {--------------------------------------------------------------------
1148 --------------------------------------------------------------------}
1152 (x:xx) -> let z' = f z x in seq z' (foldlStrict f z' xx)
1155 {--------------------------------------------------------------------
1157 --------------------------------------------------------------------}
1158 testTree :: [Int] -> IntMap Int
1159 testTree xs = fromList [(x,x*x*30696 `mod` 65521) | x <- xs]
1160 test1 = testTree [1..20]
1161 test2 = testTree [30,29..10]
1162 test3 = testTree [1,4,6,89,2323,53,43,234,5,79,12,9,24,9,8,423,8,42,4,8,9,3]
1164 {--------------------------------------------------------------------
1166 --------------------------------------------------------------------}
1171 { configMaxTest = 500
1172 , configMaxFail = 5000
1173 , configSize = \n -> (div n 2 + 3)
1174 , configEvery = \n args -> let s = show n in s ++ [ '\b' | _ <- s ]
1178 {--------------------------------------------------------------------
1179 Arbitrary, reasonably balanced trees
1180 --------------------------------------------------------------------}
1181 instance Arbitrary a => Arbitrary (IntMap a) where
1182 arbitrary = do{ ks <- arbitrary
1183 ; xs <- mapM (\k -> do{ x <- arbitrary; return (k,x)}) ks
1184 ; return (fromList xs)
1188 {--------------------------------------------------------------------
1189 Single, Insert, Delete
1190 --------------------------------------------------------------------}
1191 prop_Single :: Key -> Int -> Bool
1193 = (insert k x empty == singleton k x)
1195 prop_InsertDelete :: Key -> Int -> IntMap Int -> Property
1196 prop_InsertDelete k x t
1197 = not (member k t) ==> delete k (insert k x t) == t
1199 prop_UpdateDelete :: Key -> IntMap Int -> Bool
1200 prop_UpdateDelete k t
1201 = update (const Nothing) k t == delete k t
1204 {--------------------------------------------------------------------
1206 --------------------------------------------------------------------}
1207 prop_UnionInsert :: Key -> Int -> IntMap Int -> Bool
1208 prop_UnionInsert k x t
1209 = union (singleton k x) t == insert k x t
1211 prop_UnionAssoc :: IntMap Int -> IntMap Int -> IntMap Int -> Bool
1212 prop_UnionAssoc t1 t2 t3
1213 = union t1 (union t2 t3) == union (union t1 t2) t3
1215 prop_UnionComm :: IntMap Int -> IntMap Int -> Bool
1216 prop_UnionComm t1 t2
1217 = (union t1 t2 == unionWith (\x y -> y) t2 t1)
1220 prop_Diff :: [(Key,Int)] -> [(Key,Int)] -> Bool
1222 = List.sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys)))
1223 == List.sort ((List.\\) (nub (Prelude.map fst xs)) (nub (Prelude.map fst ys)))
1225 prop_Int :: [(Key,Int)] -> [(Key,Int)] -> Bool
1227 = List.sort (keys (intersection (fromListWith (+) xs) (fromListWith (+) ys)))
1228 == List.sort (nub ((List.intersect) (Prelude.map fst xs) (Prelude.map fst ys)))
1230 {--------------------------------------------------------------------
1232 --------------------------------------------------------------------}
1234 = forAll (choose (5,100)) $ \n ->
1235 let xs = [(x,()) | x <- [0..n::Int]]
1236 in fromAscList xs == fromList xs
1238 prop_List :: [Key] -> Bool
1240 = (sort (nub xs) == [x | (x,()) <- toAscList (fromList [(x,()) | x <- xs])])