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(..))
141 import Data.Foldable (Foldable(foldMap))
145 import qualified Prelude
146 import Debug.QuickCheck
147 import List (nub,sort)
148 import qualified List
151 #if __GLASGOW_HASKELL__
153 import Data.Generics.Basics
154 import Data.Generics.Instances
157 #if __GLASGOW_HASKELL__ >= 503
159 import GHC.Exts ( Word(..), Int(..), shiftRL# )
160 #elif __GLASGOW_HASKELL__
162 import GlaExts ( Word(..), Int(..), shiftRL# )
167 infixl 9 \\{-This comment teaches CPP correct behaviour -}
169 -- A "Nat" is a natural machine word (an unsigned Int)
172 natFromInt :: Key -> Nat
173 natFromInt i = fromIntegral i
175 intFromNat :: Nat -> Key
176 intFromNat w = fromIntegral w
178 shiftRL :: Nat -> Key -> Nat
179 #if __GLASGOW_HASKELL__
180 {--------------------------------------------------------------------
181 GHC: use unboxing to get @shiftRL@ inlined.
182 --------------------------------------------------------------------}
183 shiftRL (W# x) (I# i)
186 shiftRL x i = shiftR x i
189 {--------------------------------------------------------------------
191 --------------------------------------------------------------------}
193 -- | /O(min(n,W))/. Find the value at a key.
194 -- Calls 'error' when the element can not be found.
196 (!) :: IntMap a -> Key -> a
199 -- | /O(n+m)/. See 'difference'.
200 (\\) :: IntMap a -> IntMap b -> IntMap a
201 m1 \\ m2 = difference m1 m2
203 {--------------------------------------------------------------------
205 --------------------------------------------------------------------}
206 -- | A map of integers to values @a@.
208 | Tip {-# UNPACK #-} !Key a
209 | Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(IntMap a) !(IntMap a)
215 instance Monoid (IntMap a) where
220 instance Foldable IntMap where
221 foldMap f Nil = mempty
222 foldMap f (Tip _k v) = f v
223 foldMap f (Bin _ _ l r) = foldMap f l `mappend` foldMap f r
225 #if __GLASGOW_HASKELL__
227 {--------------------------------------------------------------------
229 --------------------------------------------------------------------}
231 -- This instance preserves data abstraction at the cost of inefficiency.
232 -- We omit reflection services for the sake of data abstraction.
234 instance Data a => Data (IntMap a) where
235 gfoldl f z im = z fromList `f` (toList im)
236 toConstr _ = error "toConstr"
237 gunfold _ _ = error "gunfold"
238 dataTypeOf _ = mkNorepType "Data.IntMap.IntMap"
239 dataCast1 f = gcast1 f
243 {--------------------------------------------------------------------
245 --------------------------------------------------------------------}
246 -- | /O(1)/. Is the map empty?
247 null :: IntMap a -> Bool
251 -- | /O(n)/. Number of elements in the map.
252 size :: IntMap a -> Int
255 Bin p m l r -> size l + size r
259 -- | /O(min(n,W))/. Is the key a member of the map?
260 member :: Key -> IntMap a -> Bool
266 -- | /O(min(n,W))/. Lookup the value at a key in the map.
267 lookup :: Key -> IntMap a -> Maybe a
269 = let nk = natFromInt k in seq nk (lookupN nk t)
271 lookupN :: Nat -> IntMap a -> Maybe a
275 | zeroN k (natFromInt m) -> lookupN k l
276 | otherwise -> lookupN k r
278 | (k == natFromInt kx) -> Just x
279 | otherwise -> Nothing
282 find' :: Key -> IntMap a -> a
285 Nothing -> error ("IntMap.find: key " ++ show k ++ " is not an element of the map")
289 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
290 -- returns the value at key @k@ or returns @def@ when the key is not an
291 -- element of the map.
292 findWithDefault :: a -> Key -> IntMap a -> a
293 findWithDefault def k m
298 {--------------------------------------------------------------------
300 --------------------------------------------------------------------}
301 -- | /O(1)/. The empty map.
306 -- | /O(1)/. A map of one element.
307 singleton :: Key -> a -> IntMap a
311 {--------------------------------------------------------------------
313 --------------------------------------------------------------------}
314 -- | /O(min(n,W))/. Insert a new key\/value pair in the map.
315 -- If the key is already present in the map, the associated value is
316 -- replaced with the supplied value, i.e. 'insert' is equivalent to
317 -- @'insertWith' 'const'@.
318 insert :: Key -> a -> IntMap a -> IntMap a
322 | nomatch k p m -> join k (Tip k x) p t
323 | zero k m -> Bin p m (insert k x l) r
324 | otherwise -> Bin p m l (insert k x r)
327 | otherwise -> join k (Tip k x) ky t
330 -- right-biased insertion, used by 'union'
331 -- | /O(min(n,W))/. Insert with a combining function.
332 -- @'insertWith' f key value mp@
333 -- will insert the pair (key, value) into @mp@ if key does
334 -- not exist in the map. If the key does exist, the function will
335 -- insert @f new_value old_value@.
336 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
338 = insertWithKey (\k x y -> f x y) k x t
340 -- | /O(min(n,W))/. Insert with a combining function.
341 -- @'insertWithKey' f key value mp@
342 -- will insert the pair (key, value) into @mp@ if key does
343 -- not exist in the map. If the key does exist, the function will
344 -- insert @f key new_value old_value@.
345 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
346 insertWithKey f k x t
349 | nomatch k p m -> join k (Tip k x) p t
350 | zero k m -> Bin p m (insertWithKey f k x l) r
351 | otherwise -> Bin p m l (insertWithKey f k x r)
353 | k==ky -> Tip k (f k x y)
354 | otherwise -> join k (Tip k x) ky t
358 -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
359 -- is a pair where the first element is equal to (@'lookup' k map@)
360 -- and the second element equal to (@'insertWithKey' f k x map@).
361 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
362 insertLookupWithKey f k x t
365 | nomatch k p m -> (Nothing,join k (Tip k x) p t)
366 | zero k m -> let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r)
367 | otherwise -> let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r')
369 | k==ky -> (Just y,Tip k (f k x y))
370 | otherwise -> (Nothing,join k (Tip k x) ky t)
371 Nil -> (Nothing,Tip k x)
374 {--------------------------------------------------------------------
376 [delete] is the inlined version of [deleteWith (\k x -> Nothing)]
377 --------------------------------------------------------------------}
378 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
379 -- a member of the map, the original map is returned.
380 delete :: Key -> IntMap a -> IntMap a
385 | zero k m -> bin p m (delete k l) r
386 | otherwise -> bin p m l (delete k r)
392 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
393 -- a member of the map, the original map is returned.
394 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
396 = adjustWithKey (\k x -> f x) k m
398 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
399 -- a member of the map, the original map is returned.
400 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
402 = updateWithKey (\k x -> Just (f k x)) k m
404 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
405 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
406 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
407 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
409 = updateWithKey (\k x -> f x) k m
411 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
412 -- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
413 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
414 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
419 | zero k m -> bin p m (updateWithKey f k l) r
420 | otherwise -> bin p m l (updateWithKey f k r)
422 | k==ky -> case (f k y) of
428 -- | /O(min(n,W))/. Lookup and update.
429 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
430 updateLookupWithKey f k t
433 | nomatch k p m -> (Nothing,t)
434 | zero k m -> let (found,l') = updateLookupWithKey f k l in (found,bin p m l' r)
435 | otherwise -> let (found,r') = updateLookupWithKey f k r in (found,bin p m l r')
437 | k==ky -> case (f k y) of
438 Just y' -> (Just y,Tip ky y')
439 Nothing -> (Just y,Nil)
440 | otherwise -> (Nothing,t)
444 {--------------------------------------------------------------------
446 --------------------------------------------------------------------}
447 -- | The union of a list of maps.
448 unions :: [IntMap a] -> IntMap a
450 = foldlStrict union empty xs
452 -- | The union of a list of maps, with a combining operation
453 unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a
455 = foldlStrict (unionWith f) empty ts
457 -- | /O(n+m)/. The (left-biased) union of two maps.
458 -- It prefers the first map when duplicate keys are encountered,
459 -- i.e. (@'union' == 'unionWith' 'const'@).
460 union :: IntMap a -> IntMap a -> IntMap a
461 union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
462 | shorter m1 m2 = union1
463 | shorter m2 m1 = union2
464 | p1 == p2 = Bin p1 m1 (union l1 l2) (union r1 r2)
465 | otherwise = join p1 t1 p2 t2
467 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
468 | zero p2 m1 = Bin p1 m1 (union l1 t2) r1
469 | otherwise = Bin p1 m1 l1 (union r1 t2)
471 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
472 | zero p1 m2 = Bin p2 m2 (union t1 l2) r2
473 | otherwise = Bin p2 m2 l2 (union t1 r2)
475 union (Tip k x) t = insert k x t
476 union t (Tip k x) = insertWith (\x y -> y) k x t -- right bias
480 -- | /O(n+m)/. The union with a combining function.
481 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
483 = unionWithKey (\k x y -> f x y) m1 m2
485 -- | /O(n+m)/. The union with a combining function.
486 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
487 unionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
488 | shorter m1 m2 = union1
489 | shorter m2 m1 = union2
490 | p1 == p2 = Bin p1 m1 (unionWithKey f l1 l2) (unionWithKey f r1 r2)
491 | otherwise = join p1 t1 p2 t2
493 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
494 | zero p2 m1 = Bin p1 m1 (unionWithKey f l1 t2) r1
495 | otherwise = Bin p1 m1 l1 (unionWithKey f r1 t2)
497 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
498 | zero p1 m2 = Bin p2 m2 (unionWithKey f t1 l2) r2
499 | otherwise = Bin p2 m2 l2 (unionWithKey f t1 r2)
501 unionWithKey f (Tip k x) t = insertWithKey f k x t
502 unionWithKey f t (Tip k x) = insertWithKey (\k x y -> f k y x) k x t -- right bias
503 unionWithKey f Nil t = t
504 unionWithKey f t Nil = t
506 {--------------------------------------------------------------------
508 --------------------------------------------------------------------}
509 -- | /O(n+m)/. Difference between two maps (based on keys).
510 difference :: IntMap a -> IntMap b -> IntMap a
511 difference t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
512 | shorter m1 m2 = difference1
513 | shorter m2 m1 = difference2
514 | p1 == p2 = bin p1 m1 (difference l1 l2) (difference r1 r2)
517 difference1 | nomatch p2 p1 m1 = t1
518 | zero p2 m1 = bin p1 m1 (difference l1 t2) r1
519 | otherwise = bin p1 m1 l1 (difference r1 t2)
521 difference2 | nomatch p1 p2 m2 = t1
522 | zero p1 m2 = difference t1 l2
523 | otherwise = difference t1 r2
525 difference t1@(Tip k x) t2
529 difference Nil t = Nil
530 difference t (Tip k x) = delete k t
533 -- | /O(n+m)/. Difference with a combining function.
534 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
535 differenceWith f m1 m2
536 = differenceWithKey (\k x y -> f x y) m1 m2
538 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
539 -- encountered, the combining function is applied to the key and both values.
540 -- If it returns 'Nothing', the element is discarded (proper set difference).
541 -- If it returns (@'Just' y@), the element is updated with a new value @y@.
542 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
543 differenceWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
544 | shorter m1 m2 = difference1
545 | shorter m2 m1 = difference2
546 | p1 == p2 = bin p1 m1 (differenceWithKey f l1 l2) (differenceWithKey f r1 r2)
549 difference1 | nomatch p2 p1 m1 = t1
550 | zero p2 m1 = bin p1 m1 (differenceWithKey f l1 t2) r1
551 | otherwise = bin p1 m1 l1 (differenceWithKey f r1 t2)
553 difference2 | nomatch p1 p2 m2 = t1
554 | zero p1 m2 = differenceWithKey f t1 l2
555 | otherwise = differenceWithKey f t1 r2
557 differenceWithKey f t1@(Tip k x) t2
558 = case lookup k t2 of
559 Just y -> case f k x y of
564 differenceWithKey f Nil t = Nil
565 differenceWithKey f t (Tip k y) = updateWithKey (\k x -> f k x y) k t
566 differenceWithKey f t Nil = t
569 {--------------------------------------------------------------------
571 --------------------------------------------------------------------}
572 -- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys).
573 intersection :: IntMap a -> IntMap b -> IntMap a
574 intersection t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
575 | shorter m1 m2 = intersection1
576 | shorter m2 m1 = intersection2
577 | p1 == p2 = bin p1 m1 (intersection l1 l2) (intersection r1 r2)
580 intersection1 | nomatch p2 p1 m1 = Nil
581 | zero p2 m1 = intersection l1 t2
582 | otherwise = intersection r1 t2
584 intersection2 | nomatch p1 p2 m2 = Nil
585 | zero p1 m2 = intersection t1 l2
586 | otherwise = intersection t1 r2
588 intersection t1@(Tip k x) t2
591 intersection t (Tip k x)
595 intersection Nil t = Nil
596 intersection t Nil = Nil
598 -- | /O(n+m)/. The intersection with a combining function.
599 intersectionWith :: (a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
600 intersectionWith f m1 m2
601 = intersectionWithKey (\k x y -> f x y) m1 m2
603 -- | /O(n+m)/. The intersection with a combining function.
604 intersectionWithKey :: (Key -> a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
605 intersectionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
606 | shorter m1 m2 = intersection1
607 | shorter m2 m1 = intersection2
608 | p1 == p2 = bin p1 m1 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2)
611 intersection1 | nomatch p2 p1 m1 = Nil
612 | zero p2 m1 = intersectionWithKey f l1 t2
613 | otherwise = intersectionWithKey f r1 t2
615 intersection2 | nomatch p1 p2 m2 = Nil
616 | zero p1 m2 = intersectionWithKey f t1 l2
617 | otherwise = intersectionWithKey f t1 r2
619 intersectionWithKey f t1@(Tip k x) t2
620 = case lookup k t2 of
621 Just y -> Tip k (f k x y)
623 intersectionWithKey f t1 (Tip k y)
624 = case lookup k t1 of
625 Just x -> Tip k (f k x y)
627 intersectionWithKey f Nil t = Nil
628 intersectionWithKey f t Nil = Nil
631 {--------------------------------------------------------------------
633 --------------------------------------------------------------------}
634 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
635 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
636 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
637 isProperSubmapOf m1 m2
638 = isProperSubmapOfBy (==) m1 m2
640 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
641 The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
642 @m1@ and @m2@ are not equal,
643 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
644 applied to their respective values. For example, the following
645 expressions are all 'True':
647 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
648 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
650 But the following are all 'False':
652 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
653 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
654 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
656 isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
657 isProperSubmapOfBy pred t1 t2
658 = case submapCmp pred t1 t2 of
662 submapCmp pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
664 | shorter m2 m1 = submapCmpLt
665 | p1 == p2 = submapCmpEq
666 | otherwise = GT -- disjoint
668 submapCmpLt | nomatch p1 p2 m2 = GT
669 | zero p1 m2 = submapCmp pred t1 l2
670 | otherwise = submapCmp pred t1 r2
671 submapCmpEq = case (submapCmp pred l1 l2, submapCmp pred r1 r2) of
677 submapCmp pred (Bin p m l r) t = GT
678 submapCmp pred (Tip kx x) (Tip ky y)
679 | (kx == ky) && pred x y = EQ
680 | otherwise = GT -- disjoint
681 submapCmp pred (Tip k x) t
683 Just y | pred x y -> LT
684 other -> GT -- disjoint
685 submapCmp pred Nil Nil = EQ
686 submapCmp pred Nil t = LT
688 -- | /O(n+m)/. Is this a submap?
689 -- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
690 isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
692 = isSubmapOfBy (==) m1 m2
695 The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if
696 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
697 applied to their respective values. For example, the following
698 expressions are all 'True':
700 > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
701 > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
702 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
704 But the following are all 'False':
706 > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)])
707 > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
708 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
711 isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
712 isSubmapOfBy pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
713 | shorter m1 m2 = False
714 | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy pred t1 l2
715 else isSubmapOfBy pred t1 r2)
716 | otherwise = (p1==p2) && isSubmapOfBy pred l1 l2 && isSubmapOfBy pred r1 r2
717 isSubmapOfBy pred (Bin p m l r) t = False
718 isSubmapOfBy pred (Tip k x) t = case lookup k t of
721 isSubmapOfBy pred Nil t = True
723 {--------------------------------------------------------------------
725 --------------------------------------------------------------------}
726 -- | /O(n)/. Map a function over all values in the map.
727 map :: (a -> b) -> IntMap a -> IntMap b
729 = mapWithKey (\k x -> f x) m
731 -- | /O(n)/. Map a function over all values in the map.
732 mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b
735 Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r)
736 Tip k x -> Tip k (f k x)
739 -- | /O(n)/. The function @'mapAccum'@ threads an accumulating
740 -- argument through the map in ascending order of keys.
741 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
743 = mapAccumWithKey (\a k x -> f a x) a m
745 -- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating
746 -- argument through the map in ascending order of keys.
747 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
748 mapAccumWithKey f a t
751 -- | /O(n)/. The function @'mapAccumL'@ threads an accumulating
752 -- argument through the map in ascending order of keys.
753 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
756 Bin p m l r -> let (a1,l') = mapAccumL f a l
757 (a2,r') = mapAccumL f a1 r
758 in (a2,Bin p m l' r')
759 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
763 -- | /O(n)/. The function @'mapAccumR'@ threads an accumulating
764 -- argument throught the map in descending order of keys.
765 mapAccumR :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
768 Bin p m l r -> let (a1,r') = mapAccumR f a r
769 (a2,l') = mapAccumR f a1 l
770 in (a2,Bin p m l' r')
771 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
774 {--------------------------------------------------------------------
776 --------------------------------------------------------------------}
777 -- | /O(n)/. Filter all values that satisfy some predicate.
778 filter :: (a -> Bool) -> IntMap a -> IntMap a
780 = filterWithKey (\k x -> p x) m
782 -- | /O(n)/. Filter all keys\/values that satisfy some predicate.
783 filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a
787 -> bin p m (filterWithKey pred l) (filterWithKey pred r)
793 -- | /O(n)/. partition the map according to some predicate. The first
794 -- map contains all elements that satisfy the predicate, the second all
795 -- elements that fail the predicate. See also 'split'.
796 partition :: (a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
798 = partitionWithKey (\k x -> p x) m
800 -- | /O(n)/. partition the map according to some predicate. The first
801 -- map contains all elements that satisfy the predicate, the second all
802 -- elements that fail the predicate. See also 'split'.
803 partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
804 partitionWithKey pred t
807 -> let (l1,l2) = partitionWithKey pred l
808 (r1,r2) = partitionWithKey pred r
809 in (bin p m l1 r1, bin p m l2 r2)
811 | pred k x -> (t,Nil)
812 | otherwise -> (Nil,t)
816 -- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@
817 -- where all keys in @map1@ are lower than @k@ and all keys in
818 -- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
819 split :: Key -> IntMap a -> (IntMap a,IntMap a)
823 | m < 0 -> (if k >= 0 -- handle negative numbers.
824 then let (lt,gt) = split' k l in (union r lt, gt)
825 else let (lt,gt) = split' k r in (lt, union gt l))
826 | otherwise -> split' k t
830 | otherwise -> (Nil,Nil)
833 split' :: Key -> IntMap a -> (IntMap a,IntMap a)
837 | nomatch k p m -> if k>p then (t,Nil) else (Nil,t)
838 | zero k m -> let (lt,gt) = split k l in (lt,union gt r)
839 | otherwise -> let (lt,gt) = split k r in (union l lt,gt)
843 | otherwise -> (Nil,Nil)
846 -- | /O(log n)/. Performs a 'split' but also returns whether the pivot
847 -- key was found in the original map.
848 splitLookup :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
852 | m < 0 -> (if k >= 0 -- handle negative numbers.
853 then let (lt,found,gt) = splitLookup' k l in (union r lt,found, gt)
854 else let (lt,found,gt) = splitLookup' k r in (lt,found, union gt l))
855 | otherwise -> splitLookup' k t
857 | k>ky -> (t,Nothing,Nil)
858 | k<ky -> (Nil,Nothing,t)
859 | otherwise -> (Nil,Just y,Nil)
860 Nil -> (Nil,Nothing,Nil)
862 splitLookup' :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
866 | nomatch k p m -> if k>p then (t,Nothing,Nil) else (Nil,Nothing,t)
867 | zero k m -> let (lt,found,gt) = splitLookup k l in (lt,found,union gt r)
868 | otherwise -> let (lt,found,gt) = splitLookup k r in (union l lt,found,gt)
870 | k>ky -> (t,Nothing,Nil)
871 | k<ky -> (Nil,Nothing,t)
872 | otherwise -> (Nil,Just y,Nil)
873 Nil -> (Nil,Nothing,Nil)
875 {--------------------------------------------------------------------
877 --------------------------------------------------------------------}
878 -- | /O(n)/. Fold the values in the map, such that
879 -- @'fold' f z == 'Prelude.foldr' f z . 'elems'@.
882 -- > elems map = fold (:) [] map
884 fold :: (a -> b -> b) -> b -> IntMap a -> b
886 = foldWithKey (\k x y -> f x y) z t
888 -- | /O(n)/. Fold the keys and values in the map, such that
889 -- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
892 -- > keys map = foldWithKey (\k x ks -> k:ks) [] map
894 foldWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
898 foldr :: (Key -> a -> b -> b) -> b -> IntMap a -> b
901 Bin 0 m l r | m < 0 -> foldr' f (foldr' f z l) r -- put negative numbers before.
902 Bin _ _ _ _ -> foldr' f z t
906 foldr' :: (Key -> a -> b -> b) -> b -> IntMap a -> b
909 Bin p m l r -> foldr' f (foldr' f z r) l
915 {--------------------------------------------------------------------
917 --------------------------------------------------------------------}
919 -- Return all elements of the map in the ascending order of their keys.
920 elems :: IntMap a -> [a]
922 = foldWithKey (\k x xs -> x:xs) [] m
924 -- | /O(n)/. Return all keys of the map in ascending order.
925 keys :: IntMap a -> [Key]
927 = foldWithKey (\k x ks -> k:ks) [] m
929 -- | /O(n*min(n,W))/. The set of all keys of the map.
930 keysSet :: IntMap a -> IntSet.IntSet
931 keysSet m = IntSet.fromDistinctAscList (keys m)
934 -- | /O(n)/. Return all key\/value pairs in the map in ascending key order.
935 assocs :: IntMap a -> [(Key,a)]
940 {--------------------------------------------------------------------
942 --------------------------------------------------------------------}
943 -- | /O(n)/. Convert the map to a list of key\/value pairs.
944 toList :: IntMap a -> [(Key,a)]
946 = foldWithKey (\k x xs -> (k,x):xs) [] t
948 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
949 -- keys are in ascending order.
950 toAscList :: IntMap a -> [(Key,a)]
952 = -- NOTE: the following algorithm only works for big-endian trees
953 let (pos,neg) = span (\(k,x) -> k >=0) (foldr (\k x xs -> (k,x):xs) [] t) in neg ++ pos
955 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
956 fromList :: [(Key,a)] -> IntMap a
958 = foldlStrict ins empty xs
960 ins t (k,x) = insert k x t
962 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
963 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
965 = fromListWithKey (\k x y -> f x y) xs
967 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
968 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
970 = foldlStrict ins empty xs
972 ins t (k,x) = insertWithKey f k x t
974 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
975 -- the keys are in ascending order.
976 fromAscList :: [(Key,a)] -> IntMap a
980 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
981 -- the keys are in ascending order, with a combining function on equal keys.
982 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
986 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
987 -- the keys are in ascending order, with a combining function on equal keys.
988 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
989 fromAscListWithKey f xs
990 = fromListWithKey f xs
992 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
993 -- the keys are in ascending order and all distinct.
994 fromDistinctAscList :: [(Key,a)] -> IntMap a
995 fromDistinctAscList xs
999 {--------------------------------------------------------------------
1001 --------------------------------------------------------------------}
1002 instance Eq a => Eq (IntMap a) where
1003 t1 == t2 = equal t1 t2
1004 t1 /= t2 = nequal t1 t2
1006 equal :: Eq a => IntMap a -> IntMap a -> Bool
1007 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1008 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
1009 equal (Tip kx x) (Tip ky y)
1010 = (kx == ky) && (x==y)
1011 equal Nil Nil = True
1014 nequal :: Eq a => IntMap a -> IntMap a -> Bool
1015 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1016 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
1017 nequal (Tip kx x) (Tip ky y)
1018 = (kx /= ky) || (x/=y)
1019 nequal Nil Nil = False
1022 {--------------------------------------------------------------------
1024 --------------------------------------------------------------------}
1026 instance Ord a => Ord (IntMap a) where
1027 compare m1 m2 = compare (toList m1) (toList m2)
1029 {--------------------------------------------------------------------
1031 --------------------------------------------------------------------}
1033 instance Functor IntMap where
1036 {--------------------------------------------------------------------
1038 --------------------------------------------------------------------}
1040 instance Show a => Show (IntMap a) where
1041 showsPrec d m = showParen (d > 10) $
1042 showString "fromList " . shows (toList m)
1044 showMap :: (Show a) => [(Key,a)] -> ShowS
1048 = showChar '{' . showElem x . showTail xs
1050 showTail [] = showChar '}'
1051 showTail (x:xs) = showChar ',' . showElem x . showTail xs
1053 showElem (k,x) = shows k . showString ":=" . shows x
1055 {--------------------------------------------------------------------
1057 --------------------------------------------------------------------}
1058 instance (Read e) => Read (IntMap e) where
1059 #ifdef __GLASGOW_HASKELL__
1060 readPrec = parens $ prec 10 $ do
1061 Ident "fromList" <- lexP
1063 return (fromList xs)
1065 readListPrec = readListPrecDefault
1067 readsPrec p = readParen (p > 10) $ \ r -> do
1068 ("fromList",s) <- lex r
1070 return (fromList xs,t)
1073 {--------------------------------------------------------------------
1075 --------------------------------------------------------------------}
1077 #include "Typeable.h"
1078 INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap")
1080 {--------------------------------------------------------------------
1082 --------------------------------------------------------------------}
1083 -- | /O(n)/. Show the tree that implements the map. The tree is shown
1084 -- in a compressed, hanging format.
1085 showTree :: Show a => IntMap a -> String
1087 = showTreeWith True False s
1090 {- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows
1091 the tree that implements the map. If @hang@ is
1092 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
1093 @wide@ is 'True', an extra wide version is shown.
1095 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
1096 showTreeWith hang wide t
1097 | hang = (showsTreeHang wide [] t) ""
1098 | otherwise = (showsTree wide [] [] t) ""
1100 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
1101 showsTree wide lbars rbars t
1104 -> showsTree wide (withBar rbars) (withEmpty rbars) r .
1105 showWide wide rbars .
1106 showsBars lbars . showString (showBin p m) . showString "\n" .
1107 showWide wide lbars .
1108 showsTree wide (withEmpty lbars) (withBar lbars) l
1110 -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1111 Nil -> showsBars lbars . showString "|\n"
1113 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
1114 showsTreeHang wide bars t
1117 -> showsBars bars . showString (showBin p m) . showString "\n" .
1118 showWide wide bars .
1119 showsTreeHang wide (withBar bars) l .
1120 showWide wide bars .
1121 showsTreeHang wide (withEmpty bars) r
1123 -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1124 Nil -> showsBars bars . showString "|\n"
1127 = "*" -- ++ show (p,m)
1130 | wide = showString (concat (reverse bars)) . showString "|\n"
1133 showsBars :: [String] -> ShowS
1137 _ -> showString (concat (reverse (tail bars))) . showString node
1140 withBar bars = "| ":bars
1141 withEmpty bars = " ":bars
1144 {--------------------------------------------------------------------
1146 --------------------------------------------------------------------}
1147 {--------------------------------------------------------------------
1149 --------------------------------------------------------------------}
1150 join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
1152 | zero p1 m = Bin p m t1 t2
1153 | otherwise = Bin p m t2 t1
1155 m = branchMask p1 p2
1158 {--------------------------------------------------------------------
1159 @bin@ assures that we never have empty trees within a tree.
1160 --------------------------------------------------------------------}
1161 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
1164 bin p m l r = Bin p m l r
1167 {--------------------------------------------------------------------
1168 Endian independent bit twiddling
1169 --------------------------------------------------------------------}
1170 zero :: Key -> Mask -> Bool
1172 = (natFromInt i) .&. (natFromInt m) == 0
1174 nomatch,match :: Key -> Prefix -> Mask -> Bool
1181 mask :: Key -> Mask -> Prefix
1183 = maskW (natFromInt i) (natFromInt m)
1186 zeroN :: Nat -> Nat -> Bool
1187 zeroN i m = (i .&. m) == 0
1189 {--------------------------------------------------------------------
1190 Big endian operations
1191 --------------------------------------------------------------------}
1192 maskW :: Nat -> Nat -> Prefix
1194 = intFromNat (i .&. (complement (m-1) `xor` m))
1196 shorter :: Mask -> Mask -> Bool
1198 = (natFromInt m1) > (natFromInt m2)
1200 branchMask :: Prefix -> Prefix -> Mask
1202 = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
1204 {----------------------------------------------------------------------
1205 Finding the highest bit (mask) in a word [x] can be done efficiently in
1207 * convert to a floating point value and the mantissa tells us the
1208 [log2(x)] that corresponds with the highest bit position. The mantissa
1209 is retrieved either via the standard C function [frexp] or by some bit
1210 twiddling on IEEE compatible numbers (float). Note that one needs to
1211 use at least [double] precision for an accurate mantissa of 32 bit
1213 * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
1214 * use processor specific assembler instruction (asm).
1216 The most portable way would be [bit], but is it efficient enough?
1217 I have measured the cycle counts of the different methods on an AMD
1218 Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
1220 highestBitMask: method cycles
1227 highestBit: method cycles
1234 Wow, the bit twiddling is on today's RISC like machines even faster
1235 than a single CISC instruction (BSR)!
1236 ----------------------------------------------------------------------}
1238 {----------------------------------------------------------------------
1239 [highestBitMask] returns a word where only the highest bit is set.
1240 It is found by first setting all bits in lower positions than the
1241 highest bit and than taking an exclusive or with the original value.
1242 Allthough the function may look expensive, GHC compiles this into
1243 excellent C code that subsequently compiled into highly efficient
1244 machine code. The algorithm is derived from Jorg Arndt's FXT library.
1245 ----------------------------------------------------------------------}
1246 highestBitMask :: Nat -> Nat
1248 = case (x .|. shiftRL x 1) of
1249 x -> case (x .|. shiftRL x 2) of
1250 x -> case (x .|. shiftRL x 4) of
1251 x -> case (x .|. shiftRL x 8) of
1252 x -> case (x .|. shiftRL x 16) of
1253 x -> case (x .|. shiftRL x 32) of -- for 64 bit platforms
1254 x -> (x `xor` (shiftRL x 1))
1257 {--------------------------------------------------------------------
1259 --------------------------------------------------------------------}
1263 (x:xx) -> let z' = f z x in seq z' (foldlStrict f z' xx)
1266 {--------------------------------------------------------------------
1268 --------------------------------------------------------------------}
1269 testTree :: [Int] -> IntMap Int
1270 testTree xs = fromList [(x,x*x*30696 `mod` 65521) | x <- xs]
1271 test1 = testTree [1..20]
1272 test2 = testTree [30,29..10]
1273 test3 = testTree [1,4,6,89,2323,53,43,234,5,79,12,9,24,9,8,423,8,42,4,8,9,3]
1275 {--------------------------------------------------------------------
1277 --------------------------------------------------------------------}
1282 { configMaxTest = 500
1283 , configMaxFail = 5000
1284 , configSize = \n -> (div n 2 + 3)
1285 , configEvery = \n args -> let s = show n in s ++ [ '\b' | _ <- s ]
1289 {--------------------------------------------------------------------
1290 Arbitrary, reasonably balanced trees
1291 --------------------------------------------------------------------}
1292 instance Arbitrary a => Arbitrary (IntMap a) where
1293 arbitrary = do{ ks <- arbitrary
1294 ; xs <- mapM (\k -> do{ x <- arbitrary; return (k,x)}) ks
1295 ; return (fromList xs)
1299 {--------------------------------------------------------------------
1300 Single, Insert, Delete
1301 --------------------------------------------------------------------}
1302 prop_Single :: Key -> Int -> Bool
1304 = (insert k x empty == singleton k x)
1306 prop_InsertDelete :: Key -> Int -> IntMap Int -> Property
1307 prop_InsertDelete k x t
1308 = not (member k t) ==> delete k (insert k x t) == t
1310 prop_UpdateDelete :: Key -> IntMap Int -> Bool
1311 prop_UpdateDelete k t
1312 = update (const Nothing) k t == delete k t
1315 {--------------------------------------------------------------------
1317 --------------------------------------------------------------------}
1318 prop_UnionInsert :: Key -> Int -> IntMap Int -> Bool
1319 prop_UnionInsert k x t
1320 = union (singleton k x) t == insert k x t
1322 prop_UnionAssoc :: IntMap Int -> IntMap Int -> IntMap Int -> Bool
1323 prop_UnionAssoc t1 t2 t3
1324 = union t1 (union t2 t3) == union (union t1 t2) t3
1326 prop_UnionComm :: IntMap Int -> IntMap Int -> Bool
1327 prop_UnionComm t1 t2
1328 = (union t1 t2 == unionWith (\x y -> y) t2 t1)
1331 prop_Diff :: [(Key,Int)] -> [(Key,Int)] -> Bool
1333 = List.sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys)))
1334 == List.sort ((List.\\) (nub (Prelude.map fst xs)) (nub (Prelude.map fst ys)))
1336 prop_Int :: [(Key,Int)] -> [(Key,Int)] -> Bool
1338 = List.sort (keys (intersection (fromListWith (+) xs) (fromListWith (+) ys)))
1339 == List.sort (nub ((List.intersect) (Prelude.map fst xs) (Prelude.map fst ys)))
1341 {--------------------------------------------------------------------
1343 --------------------------------------------------------------------}
1345 = forAll (choose (5,100)) $ \n ->
1346 let xs = [(x,()) | x <- [0..n::Int]]
1347 in fromAscList xs == fromList xs
1349 prop_List :: [Key] -> Bool
1351 = (sort (nub xs) == [x | (x,()) <- toAscList (fromList [(x,()) | x <- xs])])