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
144 import qualified Prelude
145 import Debug.QuickCheck
146 import List (nub,sort)
147 import qualified List
150 #if __GLASGOW_HASKELL__
151 import Data.Generics.Basics
152 import Data.Generics.Instances
155 #if __GLASGOW_HASKELL__ >= 503
157 import GHC.Exts ( Word(..), Int(..), shiftRL# )
158 #elif __GLASGOW_HASKELL__
160 import GlaExts ( Word(..), Int(..), shiftRL# )
165 infixl 9 \\{-This comment teaches CPP correct behaviour -}
167 -- A "Nat" is a natural machine word (an unsigned Int)
170 natFromInt :: Key -> Nat
171 natFromInt i = fromIntegral i
173 intFromNat :: Nat -> Key
174 intFromNat w = fromIntegral w
176 shiftRL :: Nat -> Key -> Nat
177 #if __GLASGOW_HASKELL__
178 {--------------------------------------------------------------------
179 GHC: use unboxing to get @shiftRL@ inlined.
180 --------------------------------------------------------------------}
181 shiftRL (W# x) (I# i)
184 shiftRL x i = shiftR x i
187 {--------------------------------------------------------------------
189 --------------------------------------------------------------------}
191 -- | /O(min(n,W))/. Find the value at a key.
192 -- Calls 'error' when the element can not be found.
194 (!) :: IntMap a -> Key -> a
197 -- | /O(n+m)/. See 'difference'.
198 (\\) :: IntMap a -> IntMap b -> IntMap a
199 m1 \\ m2 = difference m1 m2
201 {--------------------------------------------------------------------
203 --------------------------------------------------------------------}
204 -- | A map of integers to values @a@.
206 | Tip {-# UNPACK #-} !Key a
207 | Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(IntMap a) !(IntMap a)
213 #if __GLASGOW_HASKELL__
215 {--------------------------------------------------------------------
217 --------------------------------------------------------------------}
219 -- This instance preserves data abstraction at the cost of inefficiency.
220 -- We omit reflection services for the sake of data abstraction.
222 instance Data a => Data (IntMap a) where
223 gfoldl f z im = z fromList `f` (toList im)
224 toConstr _ = error "toConstr"
225 gunfold _ _ = error "gunfold"
226 dataTypeOf _ = mkNorepType "Data.IntMap.IntMap"
230 {--------------------------------------------------------------------
232 --------------------------------------------------------------------}
233 -- | /O(1)/. Is the map empty?
234 null :: IntMap a -> Bool
238 -- | /O(n)/. Number of elements in the map.
239 size :: IntMap a -> Int
242 Bin p m l r -> size l + size r
246 -- | /O(min(n,W))/. Is the key a member of the map?
247 member :: Key -> IntMap a -> Bool
253 -- | /O(min(n,W))/. Lookup the value at a key in the map.
254 lookup :: Key -> IntMap a -> Maybe a
256 = let nk = natFromInt k in seq nk (lookupN nk t)
258 lookupN :: Nat -> IntMap a -> Maybe a
262 | zeroN k (natFromInt m) -> lookupN k l
263 | otherwise -> lookupN k r
265 | (k == natFromInt kx) -> Just x
266 | otherwise -> Nothing
269 find' :: Key -> IntMap a -> a
272 Nothing -> error ("IntMap.find: key " ++ show k ++ " is not an element of the map")
276 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
277 -- returns the value at key @k@ or returns @def@ when the key is not an
278 -- element of the map.
279 findWithDefault :: a -> Key -> IntMap a -> a
280 findWithDefault def k m
285 {--------------------------------------------------------------------
287 --------------------------------------------------------------------}
288 -- | /O(1)/. The empty map.
293 -- | /O(1)/. A map of one element.
294 singleton :: Key -> a -> IntMap a
298 {--------------------------------------------------------------------
300 'insert' is the inlined version of 'insertWith (\k x y -> x)'
301 --------------------------------------------------------------------}
302 -- | /O(min(n,W))/. Insert a new key\/value pair in the map. When the key
303 -- is already an element of the set, its value is replaced by the new value,
304 -- ie. 'insert' is left-biased.
305 insert :: Key -> a -> IntMap a -> IntMap a
309 | nomatch k p m -> join k (Tip k x) p t
310 | zero k m -> Bin p m (insert k x l) r
311 | otherwise -> Bin p m l (insert k x r)
314 | otherwise -> join k (Tip k x) ky t
317 -- right-biased insertion, used by 'union'
318 -- | /O(min(n,W))/. Insert with a combining function.
319 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
321 = insertWithKey (\k x y -> f x y) k x t
323 -- | /O(min(n,W))/. Insert with a combining function.
324 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
325 insertWithKey f k x t
328 | nomatch k p m -> join k (Tip k x) p t
329 | zero k m -> Bin p m (insertWithKey f k x l) r
330 | otherwise -> Bin p m l (insertWithKey f k x r)
332 | k==ky -> Tip k (f k x y)
333 | otherwise -> join k (Tip k x) ky t
337 -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
338 -- is a pair where the first element is equal to (@'lookup' k map@)
339 -- and the second element equal to (@'insertWithKey' f k x map@).
340 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
341 insertLookupWithKey f k x t
344 | nomatch k p m -> (Nothing,join k (Tip k x) p t)
345 | zero k m -> let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r)
346 | otherwise -> let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r')
348 | k==ky -> (Just y,Tip k (f k x y))
349 | otherwise -> (Nothing,join k (Tip k x) ky t)
350 Nil -> (Nothing,Tip k x)
353 {--------------------------------------------------------------------
355 [delete] is the inlined version of [deleteWith (\k x -> Nothing)]
356 --------------------------------------------------------------------}
357 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
358 -- a member of the map, the original map is returned.
359 delete :: Key -> IntMap a -> IntMap a
364 | zero k m -> bin p m (delete k l) r
365 | otherwise -> bin p m l (delete k r)
371 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
372 -- a member of the map, the original map is returned.
373 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
375 = adjustWithKey (\k x -> f x) k m
377 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
378 -- a member of the map, the original map is returned.
379 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
381 = updateWithKey (\k x -> Just (f k x)) k m
383 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
384 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
385 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
386 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
388 = updateWithKey (\k x -> f x) k m
390 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
391 -- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
392 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
393 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
398 | zero k m -> bin p m (updateWithKey f k l) r
399 | otherwise -> bin p m l (updateWithKey f k r)
401 | k==ky -> case (f k y) of
407 -- | /O(min(n,W))/. Lookup and update.
408 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
409 updateLookupWithKey f k t
412 | nomatch k p m -> (Nothing,t)
413 | zero k m -> let (found,l') = updateLookupWithKey f k l in (found,bin p m l' r)
414 | otherwise -> let (found,r') = updateLookupWithKey f k r in (found,bin p m l r')
416 | k==ky -> case (f k y) of
417 Just y' -> (Just y,Tip ky y')
418 Nothing -> (Just y,Nil)
419 | otherwise -> (Nothing,t)
423 {--------------------------------------------------------------------
425 --------------------------------------------------------------------}
426 -- | The union of a list of maps.
427 unions :: [IntMap a] -> IntMap a
429 = foldlStrict union empty xs
431 -- | The union of a list of maps, with a combining operation
432 unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a
434 = foldlStrict (unionWith f) empty ts
436 -- | /O(n+m)/. The (left-biased) union of two sets.
437 union :: IntMap a -> IntMap a -> IntMap a
438 union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
439 | shorter m1 m2 = union1
440 | shorter m2 m1 = union2
441 | p1 == p2 = Bin p1 m1 (union l1 l2) (union r1 r2)
442 | otherwise = join p1 t1 p2 t2
444 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
445 | zero p2 m1 = Bin p1 m1 (union l1 t2) r1
446 | otherwise = Bin p1 m1 l1 (union r1 t2)
448 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
449 | zero p1 m2 = Bin p2 m2 (union t1 l2) r2
450 | otherwise = Bin p2 m2 l2 (union t1 r2)
452 union (Tip k x) t = insert k x t
453 union t (Tip k x) = insertWith (\x y -> y) k x t -- right bias
457 -- | /O(n+m)/. The union with a combining function.
458 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
460 = unionWithKey (\k x y -> f x y) m1 m2
462 -- | /O(n+m)/. The union with a combining function.
463 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
464 unionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
465 | shorter m1 m2 = union1
466 | shorter m2 m1 = union2
467 | p1 == p2 = Bin p1 m1 (unionWithKey f l1 l2) (unionWithKey f r1 r2)
468 | otherwise = join p1 t1 p2 t2
470 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
471 | zero p2 m1 = Bin p1 m1 (unionWithKey f l1 t2) r1
472 | otherwise = Bin p1 m1 l1 (unionWithKey f r1 t2)
474 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
475 | zero p1 m2 = Bin p2 m2 (unionWithKey f t1 l2) r2
476 | otherwise = Bin p2 m2 l2 (unionWithKey f t1 r2)
478 unionWithKey f (Tip k x) t = insertWithKey f k x t
479 unionWithKey f t (Tip k x) = insertWithKey (\k x y -> f k y x) k x t -- right bias
480 unionWithKey f Nil t = t
481 unionWithKey f t Nil = t
483 {--------------------------------------------------------------------
485 --------------------------------------------------------------------}
486 -- | /O(n+m)/. Difference between two maps (based on keys).
487 difference :: IntMap a -> IntMap b -> IntMap a
488 difference t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
489 | shorter m1 m2 = difference1
490 | shorter m2 m1 = difference2
491 | p1 == p2 = bin p1 m1 (difference l1 l2) (difference r1 r2)
494 difference1 | nomatch p2 p1 m1 = t1
495 | zero p2 m1 = bin p1 m1 (difference l1 t2) r1
496 | otherwise = bin p1 m1 l1 (difference r1 t2)
498 difference2 | nomatch p1 p2 m2 = t1
499 | zero p1 m2 = difference t1 l2
500 | otherwise = difference t1 r2
502 difference t1@(Tip k x) t2
506 difference Nil t = Nil
507 difference t (Tip k x) = delete k t
510 -- | /O(n+m)/. Difference with a combining function.
511 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
512 differenceWith f m1 m2
513 = differenceWithKey (\k x y -> f x y) m1 m2
515 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
516 -- encountered, the combining function is applied to the key and both values.
517 -- If it returns 'Nothing', the element is discarded (proper set difference).
518 -- If it returns (@'Just' y@), the element is updated with a new value @y@.
519 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
520 differenceWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
521 | shorter m1 m2 = difference1
522 | shorter m2 m1 = difference2
523 | p1 == p2 = bin p1 m1 (differenceWithKey f l1 l2) (differenceWithKey f r1 r2)
526 difference1 | nomatch p2 p1 m1 = t1
527 | zero p2 m1 = bin p1 m1 (differenceWithKey f l1 t2) r1
528 | otherwise = bin p1 m1 l1 (differenceWithKey f r1 t2)
530 difference2 | nomatch p1 p2 m2 = t1
531 | zero p1 m2 = differenceWithKey f t1 l2
532 | otherwise = differenceWithKey f t1 r2
534 differenceWithKey f t1@(Tip k x) t2
535 = case lookup k t2 of
536 Just y -> case f k x y of
541 differenceWithKey f Nil t = Nil
542 differenceWithKey f t (Tip k y) = updateWithKey (\k x -> f k x y) k t
543 differenceWithKey f t Nil = t
546 {--------------------------------------------------------------------
548 --------------------------------------------------------------------}
549 -- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys).
550 intersection :: IntMap a -> IntMap b -> IntMap a
551 intersection t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
552 | shorter m1 m2 = intersection1
553 | shorter m2 m1 = intersection2
554 | p1 == p2 = bin p1 m1 (intersection l1 l2) (intersection r1 r2)
557 intersection1 | nomatch p2 p1 m1 = Nil
558 | zero p2 m1 = intersection l1 t2
559 | otherwise = intersection r1 t2
561 intersection2 | nomatch p1 p2 m2 = Nil
562 | zero p1 m2 = intersection t1 l2
563 | otherwise = intersection t1 r2
565 intersection t1@(Tip k x) t2
568 intersection t (Tip k x)
572 intersection Nil t = Nil
573 intersection t Nil = Nil
575 -- | /O(n+m)/. The intersection with a combining function.
576 intersectionWith :: (a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
577 intersectionWith f m1 m2
578 = intersectionWithKey (\k x y -> f x y) m1 m2
580 -- | /O(n+m)/. The intersection with a combining function.
581 intersectionWithKey :: (Key -> a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
582 intersectionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
583 | shorter m1 m2 = intersection1
584 | shorter m2 m1 = intersection2
585 | p1 == p2 = bin p1 m1 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2)
588 intersection1 | nomatch p2 p1 m1 = Nil
589 | zero p2 m1 = intersectionWithKey f l1 t2
590 | otherwise = intersectionWithKey f r1 t2
592 intersection2 | nomatch p1 p2 m2 = Nil
593 | zero p1 m2 = intersectionWithKey f t1 l2
594 | otherwise = intersectionWithKey f t1 r2
596 intersectionWithKey f t1@(Tip k x) t2
597 = case lookup k t2 of
598 Just y -> Tip k (f k x y)
600 intersectionWithKey f t1 (Tip k y)
601 = case lookup k t1 of
602 Just x -> Tip k (f k x y)
604 intersectionWithKey f Nil t = Nil
605 intersectionWithKey f t Nil = Nil
608 {--------------------------------------------------------------------
610 --------------------------------------------------------------------}
611 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
612 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
613 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
614 isProperSubmapOf m1 m2
615 = isProperSubmapOfBy (==) m1 m2
617 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
618 The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
619 @m1@ and @m2@ are not equal,
620 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
621 applied to their respective values. For example, the following
622 expressions are all 'True':
624 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
625 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
627 But the following are all 'False':
629 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
630 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
631 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
633 isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
634 isProperSubmapOfBy pred t1 t2
635 = case submapCmp pred t1 t2 of
639 submapCmp pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
641 | shorter m2 m1 = submapCmpLt
642 | p1 == p2 = submapCmpEq
643 | otherwise = GT -- disjoint
645 submapCmpLt | nomatch p1 p2 m2 = GT
646 | zero p1 m2 = submapCmp pred t1 l2
647 | otherwise = submapCmp pred t1 r2
648 submapCmpEq = case (submapCmp pred l1 l2, submapCmp pred r1 r2) of
654 submapCmp pred (Bin p m l r) t = GT
655 submapCmp pred (Tip kx x) (Tip ky y)
656 | (kx == ky) && pred x y = EQ
657 | otherwise = GT -- disjoint
658 submapCmp pred (Tip k x) t
660 Just y | pred x y -> LT
661 other -> GT -- disjoint
662 submapCmp pred Nil Nil = EQ
663 submapCmp pred Nil t = LT
665 -- | /O(n+m)/. Is this a submap?
666 -- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
667 isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
669 = isSubmapOfBy (==) m1 m2
672 The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if
673 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
674 applied to their respective values. For example, the following
675 expressions are all 'True':
677 > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
678 > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
679 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
681 But the following are all 'False':
683 > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)])
684 > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
685 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
688 isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
689 isSubmapOfBy pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
690 | shorter m1 m2 = False
691 | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy pred t1 l2
692 else isSubmapOfBy pred t1 r2)
693 | otherwise = (p1==p2) && isSubmapOfBy pred l1 l2 && isSubmapOfBy pred r1 r2
694 isSubmapOfBy pred (Bin p m l r) t = False
695 isSubmapOfBy pred (Tip k x) t = case lookup k t of
698 isSubmapOfBy pred Nil t = True
700 {--------------------------------------------------------------------
702 --------------------------------------------------------------------}
703 -- | /O(n)/. Map a function over all values in the map.
704 map :: (a -> b) -> IntMap a -> IntMap b
706 = mapWithKey (\k x -> f x) m
708 -- | /O(n)/. Map a function over all values in the map.
709 mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b
712 Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r)
713 Tip k x -> Tip k (f k x)
716 -- | /O(n)/. The function @'mapAccum'@ threads an accumulating
717 -- argument through the map in ascending order of keys.
718 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
720 = mapAccumWithKey (\a k x -> f a x) a m
722 -- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating
723 -- argument through the map in ascending order of keys.
724 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
725 mapAccumWithKey f a t
728 -- | /O(n)/. The function @'mapAccumL'@ threads an accumulating
729 -- argument through the map in ascending order of keys.
730 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
733 Bin p m l r -> let (a1,l') = mapAccumL f a l
734 (a2,r') = mapAccumL f a1 r
735 in (a2,Bin p m l' r')
736 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
740 -- | /O(n)/. The function @'mapAccumR'@ threads an accumulating
741 -- argument throught the map in descending order of keys.
742 mapAccumR :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
745 Bin p m l r -> let (a1,r') = mapAccumR f a r
746 (a2,l') = mapAccumR f a1 l
747 in (a2,Bin p m l' r')
748 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
751 {--------------------------------------------------------------------
753 --------------------------------------------------------------------}
754 -- | /O(n)/. Filter all values that satisfy some predicate.
755 filter :: (a -> Bool) -> IntMap a -> IntMap a
757 = filterWithKey (\k x -> p x) m
759 -- | /O(n)/. Filter all keys\/values that satisfy some predicate.
760 filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a
764 -> bin p m (filterWithKey pred l) (filterWithKey pred r)
770 -- | /O(n)/. partition the map according to some predicate. The first
771 -- map contains all elements that satisfy the predicate, the second all
772 -- elements that fail the predicate. See also 'split'.
773 partition :: (a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
775 = partitionWithKey (\k x -> p x) m
777 -- | /O(n)/. partition the map according to some predicate. The first
778 -- map contains all elements that satisfy the predicate, the second all
779 -- elements that fail the predicate. See also 'split'.
780 partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
781 partitionWithKey pred t
784 -> let (l1,l2) = partitionWithKey pred l
785 (r1,r2) = partitionWithKey pred r
786 in (bin p m l1 r1, bin p m l2 r2)
788 | pred k x -> (t,Nil)
789 | otherwise -> (Nil,t)
793 -- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@
794 -- where all keys in @map1@ are lower than @k@ and all keys in
795 -- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
796 split :: Key -> IntMap a -> (IntMap a,IntMap a)
800 | zero k m -> let (lt,gt) = split k l in (lt,union gt r)
801 | otherwise -> let (lt,gt) = split k r in (union l lt,gt)
805 | otherwise -> (Nil,Nil)
808 -- | /O(log n)/. Performs a 'split' but also returns whether the pivot
809 -- key was found in the original map.
810 splitLookup :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
814 | zero k m -> let (lt,found,gt) = splitLookup k l in (lt,found,union gt r)
815 | otherwise -> let (lt,found,gt) = splitLookup k r in (union l lt,found,gt)
817 | k>ky -> (t,Nothing,Nil)
818 | k<ky -> (Nil,Nothing,t)
819 | otherwise -> (Nil,Just y,Nil)
820 Nil -> (Nil,Nothing,Nil)
822 {--------------------------------------------------------------------
824 --------------------------------------------------------------------}
825 -- | /O(n)/. Fold the values in the map, such that
826 -- @'fold' f z == 'Prelude.foldr' f z . 'elems'@.
829 -- > elems map = fold (:) [] map
831 fold :: (a -> b -> b) -> b -> IntMap a -> b
833 = foldWithKey (\k x y -> f x y) z t
835 -- | /O(n)/. Fold the keys and values in the map, such that
836 -- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
839 -- > keys map = foldWithKey (\k x ks -> k:ks) [] map
841 foldWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
845 foldr :: (Key -> a -> b -> b) -> b -> IntMap a -> b
848 Bin p m l r -> foldr f (foldr f z r) l
852 {--------------------------------------------------------------------
854 --------------------------------------------------------------------}
856 -- Return all elements of the map in the ascending order of their keys.
857 elems :: IntMap a -> [a]
859 = foldWithKey (\k x xs -> x:xs) [] m
861 -- | /O(n)/. Return all keys of the map in ascending order.
862 keys :: IntMap a -> [Key]
864 = foldWithKey (\k x ks -> k:ks) [] m
866 -- | /O(n*min(n,W))/. The set of all keys of the map.
867 keysSet :: IntMap a -> IntSet.IntSet
868 keysSet m = IntSet.fromDistinctAscList (keys m)
871 -- | /O(n)/. Return all key\/value pairs in the map in ascending key order.
872 assocs :: IntMap a -> [(Key,a)]
877 {--------------------------------------------------------------------
879 --------------------------------------------------------------------}
880 -- | /O(n)/. Convert the map to a list of key\/value pairs.
881 toList :: IntMap a -> [(Key,a)]
883 = foldWithKey (\k x xs -> (k,x):xs) [] t
885 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
886 -- keys are in ascending order.
887 toAscList :: IntMap a -> [(Key,a)]
889 = -- NOTE: the following algorithm only works for big-endian trees
890 let (pos,neg) = span (\(k,x) -> k >=0) (foldr (\k x xs -> (k,x):xs) [] t) in neg ++ pos
892 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
893 fromList :: [(Key,a)] -> IntMap a
895 = foldlStrict ins empty xs
897 ins t (k,x) = insert k x t
899 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
900 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
902 = fromListWithKey (\k x y -> f x y) xs
904 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
905 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
907 = foldlStrict ins empty xs
909 ins t (k,x) = insertWithKey f k x t
911 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
912 -- the keys are in ascending order.
913 fromAscList :: [(Key,a)] -> IntMap a
917 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
918 -- the keys are in ascending order, with a combining function on equal keys.
919 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
923 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
924 -- the keys are in ascending order, with a combining function on equal keys.
925 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
926 fromAscListWithKey f xs
927 = fromListWithKey f xs
929 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
930 -- the keys are in ascending order and all distinct.
931 fromDistinctAscList :: [(Key,a)] -> IntMap a
932 fromDistinctAscList xs
936 {--------------------------------------------------------------------
938 --------------------------------------------------------------------}
939 instance Eq a => Eq (IntMap a) where
940 t1 == t2 = equal t1 t2
941 t1 /= t2 = nequal t1 t2
943 equal :: Eq a => IntMap a -> IntMap a -> Bool
944 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
945 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
946 equal (Tip kx x) (Tip ky y)
947 = (kx == ky) && (x==y)
951 nequal :: Eq a => IntMap a -> IntMap a -> Bool
952 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
953 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
954 nequal (Tip kx x) (Tip ky y)
955 = (kx /= ky) || (x/=y)
956 nequal Nil Nil = False
959 {--------------------------------------------------------------------
961 --------------------------------------------------------------------}
963 instance Ord a => Ord (IntMap a) where
964 compare m1 m2 = compare (toList m1) (toList m2)
966 {--------------------------------------------------------------------
968 --------------------------------------------------------------------}
970 instance Functor IntMap where
973 {--------------------------------------------------------------------
975 --------------------------------------------------------------------}
977 instance Ord a => Monoid (IntMap a) where
982 {--------------------------------------------------------------------
984 --------------------------------------------------------------------}
986 instance Show a => Show (IntMap a) where
987 showsPrec d t = showMap (toList t)
990 showMap :: (Show a) => [(Key,a)] -> ShowS
994 = showChar '{' . showElem x . showTail xs
996 showTail [] = showChar '}'
997 showTail (x:xs) = showChar ',' . showElem x . showTail xs
999 showElem (k,x) = shows k . showString ":=" . shows x
1001 {--------------------------------------------------------------------
1003 --------------------------------------------------------------------}
1005 #include "Typeable.h"
1006 INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap")
1008 {--------------------------------------------------------------------
1010 --------------------------------------------------------------------}
1011 -- | /O(n)/. Show the tree that implements the map. The tree is shown
1012 -- in a compressed, hanging format.
1013 showTree :: Show a => IntMap a -> String
1015 = showTreeWith True False s
1018 {- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows
1019 the tree that implements the map. If @hang@ is
1020 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
1021 @wide@ is 'True', an extra wide version is shown.
1023 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
1024 showTreeWith hang wide t
1025 | hang = (showsTreeHang wide [] t) ""
1026 | otherwise = (showsTree wide [] [] t) ""
1028 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
1029 showsTree wide lbars rbars t
1032 -> showsTree wide (withBar rbars) (withEmpty rbars) r .
1033 showWide wide rbars .
1034 showsBars lbars . showString (showBin p m) . showString "\n" .
1035 showWide wide lbars .
1036 showsTree wide (withEmpty lbars) (withBar lbars) l
1038 -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1039 Nil -> showsBars lbars . showString "|\n"
1041 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
1042 showsTreeHang wide bars t
1045 -> showsBars bars . showString (showBin p m) . showString "\n" .
1046 showWide wide bars .
1047 showsTreeHang wide (withBar bars) l .
1048 showWide wide bars .
1049 showsTreeHang wide (withEmpty bars) r
1051 -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1052 Nil -> showsBars bars . showString "|\n"
1055 = "*" -- ++ show (p,m)
1058 | wide = showString (concat (reverse bars)) . showString "|\n"
1061 showsBars :: [String] -> ShowS
1065 _ -> showString (concat (reverse (tail bars))) . showString node
1068 withBar bars = "| ":bars
1069 withEmpty bars = " ":bars
1072 {--------------------------------------------------------------------
1074 --------------------------------------------------------------------}
1075 {--------------------------------------------------------------------
1077 --------------------------------------------------------------------}
1078 join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
1080 | zero p1 m = Bin p m t1 t2
1081 | otherwise = Bin p m t2 t1
1083 m = branchMask p1 p2
1086 {--------------------------------------------------------------------
1087 @bin@ assures that we never have empty trees within a tree.
1088 --------------------------------------------------------------------}
1089 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
1092 bin p m l r = Bin p m l r
1095 {--------------------------------------------------------------------
1096 Endian independent bit twiddling
1097 --------------------------------------------------------------------}
1098 zero :: Key -> Mask -> Bool
1100 = (natFromInt i) .&. (natFromInt m) == 0
1102 nomatch,match :: Key -> Prefix -> Mask -> Bool
1109 mask :: Key -> Mask -> Prefix
1111 = maskW (natFromInt i) (natFromInt m)
1114 zeroN :: Nat -> Nat -> Bool
1115 zeroN i m = (i .&. m) == 0
1117 {--------------------------------------------------------------------
1118 Big endian operations
1119 --------------------------------------------------------------------}
1120 maskW :: Nat -> Nat -> Prefix
1122 = intFromNat (i .&. (complement (m-1) `xor` m))
1124 shorter :: Mask -> Mask -> Bool
1126 = (natFromInt m1) > (natFromInt m2)
1128 branchMask :: Prefix -> Prefix -> Mask
1130 = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
1132 {----------------------------------------------------------------------
1133 Finding the highest bit (mask) in a word [x] can be done efficiently in
1135 * convert to a floating point value and the mantissa tells us the
1136 [log2(x)] that corresponds with the highest bit position. The mantissa
1137 is retrieved either via the standard C function [frexp] or by some bit
1138 twiddling on IEEE compatible numbers (float). Note that one needs to
1139 use at least [double] precision for an accurate mantissa of 32 bit
1141 * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
1142 * use processor specific assembler instruction (asm).
1144 The most portable way would be [bit], but is it efficient enough?
1145 I have measured the cycle counts of the different methods on an AMD
1146 Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
1148 highestBitMask: method cycles
1155 highestBit: method cycles
1162 Wow, the bit twiddling is on today's RISC like machines even faster
1163 than a single CISC instruction (BSR)!
1164 ----------------------------------------------------------------------}
1166 {----------------------------------------------------------------------
1167 [highestBitMask] returns a word where only the highest bit is set.
1168 It is found by first setting all bits in lower positions than the
1169 highest bit and than taking an exclusive or with the original value.
1170 Allthough the function may look expensive, GHC compiles this into
1171 excellent C code that subsequently compiled into highly efficient
1172 machine code. The algorithm is derived from Jorg Arndt's FXT library.
1173 ----------------------------------------------------------------------}
1174 highestBitMask :: Nat -> Nat
1176 = case (x .|. shiftRL x 1) of
1177 x -> case (x .|. shiftRL x 2) of
1178 x -> case (x .|. shiftRL x 4) of
1179 x -> case (x .|. shiftRL x 8) of
1180 x -> case (x .|. shiftRL x 16) of
1181 x -> case (x .|. shiftRL x 32) of -- for 64 bit platforms
1182 x -> (x `xor` (shiftRL x 1))
1185 {--------------------------------------------------------------------
1187 --------------------------------------------------------------------}
1191 (x:xx) -> let z' = f z x in seq z' (foldlStrict f z' xx)
1194 {--------------------------------------------------------------------
1196 --------------------------------------------------------------------}
1197 testTree :: [Int] -> IntMap Int
1198 testTree xs = fromList [(x,x*x*30696 `mod` 65521) | x <- xs]
1199 test1 = testTree [1..20]
1200 test2 = testTree [30,29..10]
1201 test3 = testTree [1,4,6,89,2323,53,43,234,5,79,12,9,24,9,8,423,8,42,4,8,9,3]
1203 {--------------------------------------------------------------------
1205 --------------------------------------------------------------------}
1210 { configMaxTest = 500
1211 , configMaxFail = 5000
1212 , configSize = \n -> (div n 2 + 3)
1213 , configEvery = \n args -> let s = show n in s ++ [ '\b' | _ <- s ]
1217 {--------------------------------------------------------------------
1218 Arbitrary, reasonably balanced trees
1219 --------------------------------------------------------------------}
1220 instance Arbitrary a => Arbitrary (IntMap a) where
1221 arbitrary = do{ ks <- arbitrary
1222 ; xs <- mapM (\k -> do{ x <- arbitrary; return (k,x)}) ks
1223 ; return (fromList xs)
1227 {--------------------------------------------------------------------
1228 Single, Insert, Delete
1229 --------------------------------------------------------------------}
1230 prop_Single :: Key -> Int -> Bool
1232 = (insert k x empty == singleton k x)
1234 prop_InsertDelete :: Key -> Int -> IntMap Int -> Property
1235 prop_InsertDelete k x t
1236 = not (member k t) ==> delete k (insert k x t) == t
1238 prop_UpdateDelete :: Key -> IntMap Int -> Bool
1239 prop_UpdateDelete k t
1240 = update (const Nothing) k t == delete k t
1243 {--------------------------------------------------------------------
1245 --------------------------------------------------------------------}
1246 prop_UnionInsert :: Key -> Int -> IntMap Int -> Bool
1247 prop_UnionInsert k x t
1248 = union (singleton k x) t == insert k x t
1250 prop_UnionAssoc :: IntMap Int -> IntMap Int -> IntMap Int -> Bool
1251 prop_UnionAssoc t1 t2 t3
1252 = union t1 (union t2 t3) == union (union t1 t2) t3
1254 prop_UnionComm :: IntMap Int -> IntMap Int -> Bool
1255 prop_UnionComm t1 t2
1256 = (union t1 t2 == unionWith (\x y -> y) t2 t1)
1259 prop_Diff :: [(Key,Int)] -> [(Key,Int)] -> Bool
1261 = List.sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys)))
1262 == List.sort ((List.\\) (nub (Prelude.map fst xs)) (nub (Prelude.map fst ys)))
1264 prop_Int :: [(Key,Int)] -> [(Key,Int)] -> Bool
1266 = List.sort (keys (intersection (fromListWith (+) xs) (fromListWith (+) ys)))
1267 == List.sort (nub ((List.intersect) (Prelude.map fst xs) (Prelude.map fst ys)))
1269 {--------------------------------------------------------------------
1271 --------------------------------------------------------------------}
1273 = forAll (choose (5,100)) $ \n ->
1274 let xs = [(x,()) | x <- [0..n::Int]]
1275 in fromAscList xs == fromList xs
1277 prop_List :: [Key] -> Bool
1279 = (sort (nub xs) == [x | (x,()) <- toAscList (fromList [(x,()) | x <- xs])])