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)
138 import qualified Data.IntSet as IntSet
139 import Data.Monoid (Monoid(..))
144 import qualified Prelude
145 import Debug.QuickCheck
146 import List (nub,sort)
147 import qualified List
150 #if __GLASGOW_HASKELL__
152 import Data.Generics.Basics
153 import Data.Generics.Instances
156 #if __GLASGOW_HASKELL__ >= 503
158 import GHC.Exts ( Word(..), Int(..), shiftRL# )
159 #elif __GLASGOW_HASKELL__
161 import GlaExts ( Word(..), Int(..), shiftRL# )
166 infixl 9 \\{-This comment teaches CPP correct behaviour -}
168 -- A "Nat" is a natural machine word (an unsigned Int)
171 natFromInt :: Key -> Nat
172 natFromInt i = fromIntegral i
174 intFromNat :: Nat -> Key
175 intFromNat w = fromIntegral w
177 shiftRL :: Nat -> Key -> Nat
178 #if __GLASGOW_HASKELL__
179 {--------------------------------------------------------------------
180 GHC: use unboxing to get @shiftRL@ inlined.
181 --------------------------------------------------------------------}
182 shiftRL (W# x) (I# i)
185 shiftRL x i = shiftR x i
188 {--------------------------------------------------------------------
190 --------------------------------------------------------------------}
192 -- | /O(min(n,W))/. Find the value at a key.
193 -- Calls 'error' when the element can not be found.
195 (!) :: IntMap a -> Key -> a
198 -- | /O(n+m)/. See 'difference'.
199 (\\) :: IntMap a -> IntMap b -> IntMap a
200 m1 \\ m2 = difference m1 m2
202 {--------------------------------------------------------------------
204 --------------------------------------------------------------------}
205 -- | A map of integers to values @a@.
207 | Tip {-# UNPACK #-} !Key a
208 | Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(IntMap a) !(IntMap a)
214 instance Ord a => Monoid (IntMap a) where
219 #if __GLASGOW_HASKELL__
221 {--------------------------------------------------------------------
223 --------------------------------------------------------------------}
225 -- This instance preserves data abstraction at the cost of inefficiency.
226 -- We omit reflection services for the sake of data abstraction.
228 instance Data a => Data (IntMap a) where
229 gfoldl f z im = z fromList `f` (toList im)
230 toConstr _ = error "toConstr"
231 gunfold _ _ = error "gunfold"
232 dataTypeOf _ = mkNorepType "Data.IntMap.IntMap"
236 {--------------------------------------------------------------------
238 --------------------------------------------------------------------}
239 -- | /O(1)/. Is the map empty?
240 null :: IntMap a -> Bool
244 -- | /O(n)/. Number of elements in the map.
245 size :: IntMap a -> Int
248 Bin p m l r -> size l + size r
252 -- | /O(min(n,W))/. Is the key a member of the map?
253 member :: Key -> IntMap a -> Bool
259 -- | /O(min(n,W))/. Lookup the value at a key in the map.
260 lookup :: Key -> IntMap a -> Maybe a
262 = let nk = natFromInt k in seq nk (lookupN nk t)
264 lookupN :: Nat -> IntMap a -> Maybe a
268 | zeroN k (natFromInt m) -> lookupN k l
269 | otherwise -> lookupN k r
271 | (k == natFromInt kx) -> Just x
272 | otherwise -> Nothing
275 find' :: Key -> IntMap a -> a
278 Nothing -> error ("IntMap.find: key " ++ show k ++ " is not an element of the map")
282 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
283 -- returns the value at key @k@ or returns @def@ when the key is not an
284 -- element of the map.
285 findWithDefault :: a -> Key -> IntMap a -> a
286 findWithDefault def k m
291 {--------------------------------------------------------------------
293 --------------------------------------------------------------------}
294 -- | /O(1)/. The empty map.
299 -- | /O(1)/. A map of one element.
300 singleton :: Key -> a -> IntMap a
304 {--------------------------------------------------------------------
306 --------------------------------------------------------------------}
307 -- | /O(min(n,W))/. Insert a new key\/value pair in the map.
308 -- If the key is already present in the map, the associated value is
309 -- replaced with the supplied value, i.e. 'insert' is equivalent to
310 -- @'insertWith' 'const'@.
311 insert :: Key -> a -> IntMap a -> IntMap a
315 | nomatch k p m -> join k (Tip k x) p t
316 | zero k m -> Bin p m (insert k x l) r
317 | otherwise -> Bin p m l (insert k x r)
320 | otherwise -> join k (Tip k x) ky t
323 -- right-biased insertion, used by 'union'
324 -- | /O(min(n,W))/. Insert with a combining function.
325 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
327 = insertWithKey (\k x y -> f x y) k x t
329 -- | /O(min(n,W))/. Insert with a combining function.
330 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
331 insertWithKey f k x t
334 | nomatch k p m -> join k (Tip k x) p t
335 | zero k m -> Bin p m (insertWithKey f k x l) r
336 | otherwise -> Bin p m l (insertWithKey f k x r)
338 | k==ky -> Tip k (f k x y)
339 | otherwise -> join k (Tip k x) ky t
343 -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
344 -- is a pair where the first element is equal to (@'lookup' k map@)
345 -- and the second element equal to (@'insertWithKey' f k x map@).
346 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
347 insertLookupWithKey f k x t
350 | nomatch k p m -> (Nothing,join k (Tip k x) p t)
351 | zero k m -> let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r)
352 | otherwise -> let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r')
354 | k==ky -> (Just y,Tip k (f k x y))
355 | otherwise -> (Nothing,join k (Tip k x) ky t)
356 Nil -> (Nothing,Tip k x)
359 {--------------------------------------------------------------------
361 [delete] is the inlined version of [deleteWith (\k x -> Nothing)]
362 --------------------------------------------------------------------}
363 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
364 -- a member of the map, the original map is returned.
365 delete :: Key -> IntMap a -> IntMap a
370 | zero k m -> bin p m (delete k l) r
371 | otherwise -> bin p m l (delete k r)
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 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
381 = adjustWithKey (\k x -> f x) k m
383 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
384 -- a member of the map, the original map is returned.
385 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
387 = updateWithKey (\k x -> Just (f k x)) k m
389 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
390 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
391 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
392 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
394 = updateWithKey (\k x -> f x) k m
396 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
397 -- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
398 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
399 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
404 | zero k m -> bin p m (updateWithKey f k l) r
405 | otherwise -> bin p m l (updateWithKey f k r)
407 | k==ky -> case (f k y) of
413 -- | /O(min(n,W))/. Lookup and update.
414 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
415 updateLookupWithKey f k t
418 | nomatch k p m -> (Nothing,t)
419 | zero k m -> let (found,l') = updateLookupWithKey f k l in (found,bin p m l' r)
420 | otherwise -> let (found,r') = updateLookupWithKey f k r in (found,bin p m l r')
422 | k==ky -> case (f k y) of
423 Just y' -> (Just y,Tip ky y')
424 Nothing -> (Just y,Nil)
425 | otherwise -> (Nothing,t)
429 {--------------------------------------------------------------------
431 --------------------------------------------------------------------}
432 -- | The union of a list of maps.
433 unions :: [IntMap a] -> IntMap a
435 = foldlStrict union empty xs
437 -- | The union of a list of maps, with a combining operation
438 unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a
440 = foldlStrict (unionWith f) empty ts
442 -- | /O(n+m)/. The (left-biased) union of two maps.
443 -- It prefers the first map when duplicate keys are encountered,
444 -- i.e. (@'union' == 'unionWith' 'const'@).
445 union :: IntMap a -> IntMap a -> IntMap a
446 union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
447 | shorter m1 m2 = union1
448 | shorter m2 m1 = union2
449 | p1 == p2 = Bin p1 m1 (union l1 l2) (union r1 r2)
450 | otherwise = join p1 t1 p2 t2
452 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
453 | zero p2 m1 = Bin p1 m1 (union l1 t2) r1
454 | otherwise = Bin p1 m1 l1 (union r1 t2)
456 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
457 | zero p1 m2 = Bin p2 m2 (union t1 l2) r2
458 | otherwise = Bin p2 m2 l2 (union t1 r2)
460 union (Tip k x) t = insert k x t
461 union t (Tip k x) = insertWith (\x y -> y) k x t -- right bias
465 -- | /O(n+m)/. The union with a combining function.
466 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
468 = unionWithKey (\k x y -> f x y) m1 m2
470 -- | /O(n+m)/. The union with a combining function.
471 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
472 unionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
473 | shorter m1 m2 = union1
474 | shorter m2 m1 = union2
475 | p1 == p2 = Bin p1 m1 (unionWithKey f l1 l2) (unionWithKey f r1 r2)
476 | otherwise = join p1 t1 p2 t2
478 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
479 | zero p2 m1 = Bin p1 m1 (unionWithKey f l1 t2) r1
480 | otherwise = Bin p1 m1 l1 (unionWithKey f r1 t2)
482 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
483 | zero p1 m2 = Bin p2 m2 (unionWithKey f t1 l2) r2
484 | otherwise = Bin p2 m2 l2 (unionWithKey f t1 r2)
486 unionWithKey f (Tip k x) t = insertWithKey f k x t
487 unionWithKey f t (Tip k x) = insertWithKey (\k x y -> f k y x) k x t -- right bias
488 unionWithKey f Nil t = t
489 unionWithKey f t Nil = t
491 {--------------------------------------------------------------------
493 --------------------------------------------------------------------}
494 -- | /O(n+m)/. Difference between two maps (based on keys).
495 difference :: IntMap a -> IntMap b -> IntMap a
496 difference t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
497 | shorter m1 m2 = difference1
498 | shorter m2 m1 = difference2
499 | p1 == p2 = bin p1 m1 (difference l1 l2) (difference r1 r2)
502 difference1 | nomatch p2 p1 m1 = t1
503 | zero p2 m1 = bin p1 m1 (difference l1 t2) r1
504 | otherwise = bin p1 m1 l1 (difference r1 t2)
506 difference2 | nomatch p1 p2 m2 = t1
507 | zero p1 m2 = difference t1 l2
508 | otherwise = difference t1 r2
510 difference t1@(Tip k x) t2
514 difference Nil t = Nil
515 difference t (Tip k x) = delete k t
518 -- | /O(n+m)/. Difference with a combining function.
519 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
520 differenceWith f m1 m2
521 = differenceWithKey (\k x y -> f x y) m1 m2
523 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
524 -- encountered, the combining function is applied to the key and both values.
525 -- If it returns 'Nothing', the element is discarded (proper set difference).
526 -- If it returns (@'Just' y@), the element is updated with a new value @y@.
527 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
528 differenceWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
529 | shorter m1 m2 = difference1
530 | shorter m2 m1 = difference2
531 | p1 == p2 = bin p1 m1 (differenceWithKey f l1 l2) (differenceWithKey f r1 r2)
534 difference1 | nomatch p2 p1 m1 = t1
535 | zero p2 m1 = bin p1 m1 (differenceWithKey f l1 t2) r1
536 | otherwise = bin p1 m1 l1 (differenceWithKey f r1 t2)
538 difference2 | nomatch p1 p2 m2 = t1
539 | zero p1 m2 = differenceWithKey f t1 l2
540 | otherwise = differenceWithKey f t1 r2
542 differenceWithKey f t1@(Tip k x) t2
543 = case lookup k t2 of
544 Just y -> case f k x y of
549 differenceWithKey f Nil t = Nil
550 differenceWithKey f t (Tip k y) = updateWithKey (\k x -> f k x y) k t
551 differenceWithKey f t Nil = t
554 {--------------------------------------------------------------------
556 --------------------------------------------------------------------}
557 -- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys).
558 intersection :: IntMap a -> IntMap b -> IntMap a
559 intersection t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
560 | shorter m1 m2 = intersection1
561 | shorter m2 m1 = intersection2
562 | p1 == p2 = bin p1 m1 (intersection l1 l2) (intersection r1 r2)
565 intersection1 | nomatch p2 p1 m1 = Nil
566 | zero p2 m1 = intersection l1 t2
567 | otherwise = intersection r1 t2
569 intersection2 | nomatch p1 p2 m2 = Nil
570 | zero p1 m2 = intersection t1 l2
571 | otherwise = intersection t1 r2
573 intersection t1@(Tip k x) t2
576 intersection t (Tip k x)
580 intersection Nil t = Nil
581 intersection t Nil = Nil
583 -- | /O(n+m)/. The intersection with a combining function.
584 intersectionWith :: (a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
585 intersectionWith f m1 m2
586 = intersectionWithKey (\k x y -> f x y) m1 m2
588 -- | /O(n+m)/. The intersection with a combining function.
589 intersectionWithKey :: (Key -> a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
590 intersectionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
591 | shorter m1 m2 = intersection1
592 | shorter m2 m1 = intersection2
593 | p1 == p2 = bin p1 m1 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2)
596 intersection1 | nomatch p2 p1 m1 = Nil
597 | zero p2 m1 = intersectionWithKey f l1 t2
598 | otherwise = intersectionWithKey f r1 t2
600 intersection2 | nomatch p1 p2 m2 = Nil
601 | zero p1 m2 = intersectionWithKey f t1 l2
602 | otherwise = intersectionWithKey f t1 r2
604 intersectionWithKey f t1@(Tip k x) t2
605 = case lookup k t2 of
606 Just y -> Tip k (f k x y)
608 intersectionWithKey f t1 (Tip k y)
609 = case lookup k t1 of
610 Just x -> Tip k (f k x y)
612 intersectionWithKey f Nil t = Nil
613 intersectionWithKey f t Nil = Nil
616 {--------------------------------------------------------------------
618 --------------------------------------------------------------------}
619 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
620 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
621 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
622 isProperSubmapOf m1 m2
623 = isProperSubmapOfBy (==) m1 m2
625 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
626 The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
627 @m1@ and @m2@ are not equal,
628 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
629 applied to their respective values. For example, the following
630 expressions are all 'True':
632 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
633 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
635 But the following are all 'False':
637 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
638 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
639 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
641 isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
642 isProperSubmapOfBy pred t1 t2
643 = case submapCmp pred t1 t2 of
647 submapCmp pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
649 | shorter m2 m1 = submapCmpLt
650 | p1 == p2 = submapCmpEq
651 | otherwise = GT -- disjoint
653 submapCmpLt | nomatch p1 p2 m2 = GT
654 | zero p1 m2 = submapCmp pred t1 l2
655 | otherwise = submapCmp pred t1 r2
656 submapCmpEq = case (submapCmp pred l1 l2, submapCmp pred r1 r2) of
662 submapCmp pred (Bin p m l r) t = GT
663 submapCmp pred (Tip kx x) (Tip ky y)
664 | (kx == ky) && pred x y = EQ
665 | otherwise = GT -- disjoint
666 submapCmp pred (Tip k x) t
668 Just y | pred x y -> LT
669 other -> GT -- disjoint
670 submapCmp pred Nil Nil = EQ
671 submapCmp pred Nil t = LT
673 -- | /O(n+m)/. Is this a submap?
674 -- 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 ascending order of keys.
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 ascending order of keys.
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 ascending order of keys.
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 descending order of keys.
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 | nomatch k p m -> if k>p then (t,Nil) else (Nil,t)
809 | zero k m -> let (lt,gt) = split k l in (lt,union gt r)
810 | otherwise -> let (lt,gt) = split k r in (union l lt,gt)
814 | otherwise -> (Nil,Nil)
817 -- | /O(log n)/. Performs a 'split' but also returns whether the pivot
818 -- key was found in the original map.
819 splitLookup :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
823 | nomatch k p m -> if k>p then (t,Nothing,Nil) else (Nil,Nothing,t)
824 | zero k m -> let (lt,found,gt) = splitLookup k l in (lt,found,union gt r)
825 | otherwise -> let (lt,found,gt) = splitLookup k r in (union l lt,found,gt)
827 | k>ky -> (t,Nothing,Nil)
828 | k<ky -> (Nil,Nothing,t)
829 | otherwise -> (Nil,Just y,Nil)
830 Nil -> (Nil,Nothing,Nil)
832 {--------------------------------------------------------------------
834 --------------------------------------------------------------------}
835 -- | /O(n)/. Fold the values in the map, such that
836 -- @'fold' f z == 'Prelude.foldr' f z . 'elems'@.
839 -- > elems map = fold (:) [] map
841 fold :: (a -> b -> b) -> b -> IntMap a -> b
843 = foldWithKey (\k x y -> f x y) z t
845 -- | /O(n)/. Fold the keys and values in the map, such that
846 -- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
849 -- > keys map = foldWithKey (\k x ks -> k:ks) [] map
851 foldWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
855 foldr :: (Key -> a -> b -> b) -> b -> IntMap a -> b
858 Bin p m l r -> foldr f (foldr f z r) l
862 {--------------------------------------------------------------------
864 --------------------------------------------------------------------}
866 -- Return all elements of the map in the ascending order of their keys.
867 elems :: IntMap a -> [a]
869 = foldWithKey (\k x xs -> x:xs) [] m
871 -- | /O(n)/. Return all keys of the map in ascending order.
872 keys :: IntMap a -> [Key]
874 = foldWithKey (\k x ks -> k:ks) [] m
876 -- | /O(n*min(n,W))/. The set of all keys of the map.
877 keysSet :: IntMap a -> IntSet.IntSet
878 keysSet m = IntSet.fromDistinctAscList (keys m)
881 -- | /O(n)/. Return all key\/value pairs in the map in ascending key order.
882 assocs :: IntMap a -> [(Key,a)]
887 {--------------------------------------------------------------------
889 --------------------------------------------------------------------}
890 -- | /O(n)/. Convert the map to a list of key\/value pairs.
891 toList :: IntMap a -> [(Key,a)]
893 = foldWithKey (\k x xs -> (k,x):xs) [] t
895 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
896 -- keys are in ascending order.
897 toAscList :: IntMap a -> [(Key,a)]
899 = -- NOTE: the following algorithm only works for big-endian trees
900 let (pos,neg) = span (\(k,x) -> k >=0) (foldr (\k x xs -> (k,x):xs) [] t) in neg ++ pos
902 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
903 fromList :: [(Key,a)] -> IntMap a
905 = foldlStrict ins empty xs
907 ins t (k,x) = insert k x t
909 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
910 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
912 = fromListWithKey (\k x y -> f x y) xs
914 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
915 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
917 = foldlStrict ins empty xs
919 ins t (k,x) = insertWithKey f k x t
921 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
922 -- the keys are in ascending order.
923 fromAscList :: [(Key,a)] -> IntMap a
927 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
928 -- the keys are in ascending order, with a combining function on equal keys.
929 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
933 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
934 -- the keys are in ascending order, with a combining function on equal keys.
935 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
936 fromAscListWithKey f xs
937 = fromListWithKey f xs
939 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
940 -- the keys are in ascending order and all distinct.
941 fromDistinctAscList :: [(Key,a)] -> IntMap a
942 fromDistinctAscList xs
946 {--------------------------------------------------------------------
948 --------------------------------------------------------------------}
949 instance Eq a => Eq (IntMap a) where
950 t1 == t2 = equal t1 t2
951 t1 /= t2 = nequal t1 t2
953 equal :: Eq a => IntMap a -> IntMap a -> Bool
954 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
955 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
956 equal (Tip kx x) (Tip ky y)
957 = (kx == ky) && (x==y)
961 nequal :: Eq a => IntMap a -> IntMap a -> Bool
962 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
963 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
964 nequal (Tip kx x) (Tip ky y)
965 = (kx /= ky) || (x/=y)
966 nequal Nil Nil = False
969 {--------------------------------------------------------------------
971 --------------------------------------------------------------------}
973 instance Ord a => Ord (IntMap a) where
974 compare m1 m2 = compare (toList m1) (toList m2)
976 {--------------------------------------------------------------------
978 --------------------------------------------------------------------}
980 instance Functor IntMap where
983 {--------------------------------------------------------------------
985 --------------------------------------------------------------------}
987 instance Show a => Show (IntMap a) where
988 showsPrec d m = showParen (d > 10) $
989 showString "fromList " . shows (toList m)
991 showMap :: (Show a) => [(Key,a)] -> ShowS
995 = showChar '{' . showElem x . showTail xs
997 showTail [] = showChar '}'
998 showTail (x:xs) = showChar ',' . showElem x . showTail xs
1000 showElem (k,x) = shows k . showString ":=" . shows x
1002 {--------------------------------------------------------------------
1004 --------------------------------------------------------------------}
1005 instance (Read e) => Read (IntMap e) where
1006 #ifdef __GLASGOW_HASKELL__
1007 readPrec = parens $ prec 10 $ do
1008 Ident "fromList" <- lexP
1010 return (fromList xs)
1012 readListPrec = readListPrecDefault
1014 readsPrec p = readParen (p > 10) $ \ r -> do
1015 ("fromList",s) <- lex r
1017 return (fromList xs,t)
1020 {--------------------------------------------------------------------
1022 --------------------------------------------------------------------}
1024 #include "Typeable.h"
1025 INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap")
1027 {--------------------------------------------------------------------
1029 --------------------------------------------------------------------}
1030 -- | /O(n)/. Show the tree that implements the map. The tree is shown
1031 -- in a compressed, hanging format.
1032 showTree :: Show a => IntMap a -> String
1034 = showTreeWith True False s
1037 {- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows
1038 the tree that implements the map. If @hang@ is
1039 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
1040 @wide@ is 'True', an extra wide version is shown.
1042 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
1043 showTreeWith hang wide t
1044 | hang = (showsTreeHang wide [] t) ""
1045 | otherwise = (showsTree wide [] [] t) ""
1047 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
1048 showsTree wide lbars rbars t
1051 -> showsTree wide (withBar rbars) (withEmpty rbars) r .
1052 showWide wide rbars .
1053 showsBars lbars . showString (showBin p m) . showString "\n" .
1054 showWide wide lbars .
1055 showsTree wide (withEmpty lbars) (withBar lbars) l
1057 -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1058 Nil -> showsBars lbars . showString "|\n"
1060 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
1061 showsTreeHang wide bars t
1064 -> showsBars bars . showString (showBin p m) . showString "\n" .
1065 showWide wide bars .
1066 showsTreeHang wide (withBar bars) l .
1067 showWide wide bars .
1068 showsTreeHang wide (withEmpty bars) r
1070 -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1071 Nil -> showsBars bars . showString "|\n"
1074 = "*" -- ++ show (p,m)
1077 | wide = showString (concat (reverse bars)) . showString "|\n"
1080 showsBars :: [String] -> ShowS
1084 _ -> showString (concat (reverse (tail bars))) . showString node
1087 withBar bars = "| ":bars
1088 withEmpty bars = " ":bars
1091 {--------------------------------------------------------------------
1093 --------------------------------------------------------------------}
1094 {--------------------------------------------------------------------
1096 --------------------------------------------------------------------}
1097 join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
1099 | zero p1 m = Bin p m t1 t2
1100 | otherwise = Bin p m t2 t1
1102 m = branchMask p1 p2
1105 {--------------------------------------------------------------------
1106 @bin@ assures that we never have empty trees within a tree.
1107 --------------------------------------------------------------------}
1108 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
1111 bin p m l r = Bin p m l r
1114 {--------------------------------------------------------------------
1115 Endian independent bit twiddling
1116 --------------------------------------------------------------------}
1117 zero :: Key -> Mask -> Bool
1119 = (natFromInt i) .&. (natFromInt m) == 0
1121 nomatch,match :: Key -> Prefix -> Mask -> Bool
1128 mask :: Key -> Mask -> Prefix
1130 = maskW (natFromInt i) (natFromInt m)
1133 zeroN :: Nat -> Nat -> Bool
1134 zeroN i m = (i .&. m) == 0
1136 {--------------------------------------------------------------------
1137 Big endian operations
1138 --------------------------------------------------------------------}
1139 maskW :: Nat -> Nat -> Prefix
1141 = intFromNat (i .&. (complement (m-1) `xor` m))
1143 shorter :: Mask -> Mask -> Bool
1145 = (natFromInt m1) > (natFromInt m2)
1147 branchMask :: Prefix -> Prefix -> Mask
1149 = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
1151 {----------------------------------------------------------------------
1152 Finding the highest bit (mask) in a word [x] can be done efficiently in
1154 * convert to a floating point value and the mantissa tells us the
1155 [log2(x)] that corresponds with the highest bit position. The mantissa
1156 is retrieved either via the standard C function [frexp] or by some bit
1157 twiddling on IEEE compatible numbers (float). Note that one needs to
1158 use at least [double] precision for an accurate mantissa of 32 bit
1160 * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
1161 * use processor specific assembler instruction (asm).
1163 The most portable way would be [bit], but is it efficient enough?
1164 I have measured the cycle counts of the different methods on an AMD
1165 Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
1167 highestBitMask: method cycles
1174 highestBit: method cycles
1181 Wow, the bit twiddling is on today's RISC like machines even faster
1182 than a single CISC instruction (BSR)!
1183 ----------------------------------------------------------------------}
1185 {----------------------------------------------------------------------
1186 [highestBitMask] returns a word where only the highest bit is set.
1187 It is found by first setting all bits in lower positions than the
1188 highest bit and than taking an exclusive or with the original value.
1189 Allthough the function may look expensive, GHC compiles this into
1190 excellent C code that subsequently compiled into highly efficient
1191 machine code. The algorithm is derived from Jorg Arndt's FXT library.
1192 ----------------------------------------------------------------------}
1193 highestBitMask :: Nat -> Nat
1195 = case (x .|. shiftRL x 1) of
1196 x -> case (x .|. shiftRL x 2) of
1197 x -> case (x .|. shiftRL x 4) of
1198 x -> case (x .|. shiftRL x 8) of
1199 x -> case (x .|. shiftRL x 16) of
1200 x -> case (x .|. shiftRL x 32) of -- for 64 bit platforms
1201 x -> (x `xor` (shiftRL x 1))
1204 {--------------------------------------------------------------------
1206 --------------------------------------------------------------------}
1210 (x:xx) -> let z' = f z x in seq z' (foldlStrict f z' xx)
1213 {--------------------------------------------------------------------
1215 --------------------------------------------------------------------}
1216 testTree :: [Int] -> IntMap Int
1217 testTree xs = fromList [(x,x*x*30696 `mod` 65521) | x <- xs]
1218 test1 = testTree [1..20]
1219 test2 = testTree [30,29..10]
1220 test3 = testTree [1,4,6,89,2323,53,43,234,5,79,12,9,24,9,8,423,8,42,4,8,9,3]
1222 {--------------------------------------------------------------------
1224 --------------------------------------------------------------------}
1229 { configMaxTest = 500
1230 , configMaxFail = 5000
1231 , configSize = \n -> (div n 2 + 3)
1232 , configEvery = \n args -> let s = show n in s ++ [ '\b' | _ <- s ]
1236 {--------------------------------------------------------------------
1237 Arbitrary, reasonably balanced trees
1238 --------------------------------------------------------------------}
1239 instance Arbitrary a => Arbitrary (IntMap a) where
1240 arbitrary = do{ ks <- arbitrary
1241 ; xs <- mapM (\k -> do{ x <- arbitrary; return (k,x)}) ks
1242 ; return (fromList xs)
1246 {--------------------------------------------------------------------
1247 Single, Insert, Delete
1248 --------------------------------------------------------------------}
1249 prop_Single :: Key -> Int -> Bool
1251 = (insert k x empty == singleton k x)
1253 prop_InsertDelete :: Key -> Int -> IntMap Int -> Property
1254 prop_InsertDelete k x t
1255 = not (member k t) ==> delete k (insert k x t) == t
1257 prop_UpdateDelete :: Key -> IntMap Int -> Bool
1258 prop_UpdateDelete k t
1259 = update (const Nothing) k t == delete k t
1262 {--------------------------------------------------------------------
1264 --------------------------------------------------------------------}
1265 prop_UnionInsert :: Key -> Int -> IntMap Int -> Bool
1266 prop_UnionInsert k x t
1267 = union (singleton k x) t == insert k x t
1269 prop_UnionAssoc :: IntMap Int -> IntMap Int -> IntMap Int -> Bool
1270 prop_UnionAssoc t1 t2 t3
1271 = union t1 (union t2 t3) == union (union t1 t2) t3
1273 prop_UnionComm :: IntMap Int -> IntMap Int -> Bool
1274 prop_UnionComm t1 t2
1275 = (union t1 t2 == unionWith (\x y -> y) t2 t1)
1278 prop_Diff :: [(Key,Int)] -> [(Key,Int)] -> Bool
1280 = List.sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys)))
1281 == List.sort ((List.\\) (nub (Prelude.map fst xs)) (nub (Prelude.map fst ys)))
1283 prop_Int :: [(Key,Int)] -> [(Key,Int)] -> Bool
1285 = List.sort (keys (intersection (fromListWith (+) xs) (fromListWith (+) ys)))
1286 == List.sort (nub ((List.intersect) (Prelude.map fst xs) (Prelude.map fst ys)))
1288 {--------------------------------------------------------------------
1290 --------------------------------------------------------------------}
1292 = forAll (choose (5,100)) $ \n ->
1293 let xs = [(x,()) | x <- [0..n::Int]]
1294 in fromAscList xs == fromList xs
1296 prop_List :: [Key] -> Bool
1298 = (sort (nub xs) == [x | (x,()) <- toAscList (fromList [(x,()) | x <- xs])])