99b2ddbac3dd846cec2e0889f2f91877d1e5c2af
[packages/containers.git] / Data / IntMap.hs
1 {-# LANGUAGE CPP, NoBangPatterns, MagicHash, ScopedTypeVariables #-}
2 -----------------------------------------------------------------------------
3 -- |
4 -- Module : Data.IntMap
5 -- Copyright : (c) Daan Leijen 2002
6 -- (c) Andriy Palamarchuk 2008
7 -- License : BSD-style
8 -- Maintainer : libraries@haskell.org
9 -- Stability : provisional
10 -- Portability : portable
11 --
12 -- An efficient implementation of maps from integer keys to values.
13 --
14 -- Since many function names (but not the type name) clash with
15 -- "Prelude" names, this module is usually imported @qualified@, e.g.
16 --
17 -- > import Data.IntMap (IntMap)
18 -- > import qualified Data.IntMap as IntMap
19 --
20 -- The implementation is based on /big-endian patricia trees/. This data
21 -- structure performs especially well on binary operations like 'union'
22 -- and 'intersection'. However, my benchmarks show that it is also
23 -- (much) faster on insertions and deletions when compared to a generic
24 -- size-balanced map implementation (see "Data.Map").
25 --
26 -- * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\",
27 -- Workshop on ML, September 1998, pages 77-86,
28 -- <http://citeseer.ist.psu.edu/okasaki98fast.html>
29 --
30 -- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve
31 -- Information Coded In Alphanumeric/\", Journal of the ACM, 15(4),
32 -- October 1968, pages 514-534.
33 --
34 -- Operation comments contain the operation time complexity in
35 -- the Big-O notation <http://en.wikipedia.org/wiki/Big_O_notation>.
36 -- Many operations have a worst-case complexity of /O(min(n,W))/.
37 -- This means that the operation can become linear in the number of
38 -- elements with a maximum of /W/ -- the number of bits in an 'Int'
39 -- (32 or 64).
40 -----------------------------------------------------------------------------
41
42 -- It is essential that the bit fiddling functions like mask, zero, branchMask
43 -- etc are inlined. If they do not, the memory allocation skyrockets. The GHC
44 -- usually gets it right, but it is disastrous if it does not. Therefore we
45 -- explicitly mark these functions INLINE.
46
47 module Data.IntMap (
48 -- * Map type
49 #if !defined(TESTING)
50 IntMap, Key -- instance Eq,Show
51 #else
52 IntMap(..), Key -- instance Eq,Show
53 #endif
54
55 -- * Operators
56 , (!), (\\)
57
58 -- * Query
59 , null
60 , size
61 , member
62 , notMember
63 , lookup
64 , findWithDefault
65
66 -- * Construction
67 , empty
68 , singleton
69
70 -- ** Insertion
71 , insert
72 , insertWith
73 , insertWith'
74 , insertWithKey
75 , insertWithKey'
76 , insertLookupWithKey
77
78 -- ** Delete\/Update
79 , delete
80 , adjust
81 , adjustWithKey
82 , update
83 , updateWithKey
84 , updateLookupWithKey
85 , alter
86
87 -- * Combine
88
89 -- ** Union
90 , union
91 , unionWith
92 , unionWithKey
93 , unions
94 , unionsWith
95
96 -- ** Difference
97 , difference
98 , differenceWith
99 , differenceWithKey
100
101 -- ** Intersection
102 , intersection
103 , intersectionWith
104 , intersectionWithKey
105
106 -- * Traversal
107 -- ** Map
108 , map
109 , mapWithKey
110 , mapAccum
111 , mapAccumWithKey
112 , mapAccumRWithKey
113
114 -- * Folds
115 , foldr
116 , foldl
117 , foldrWithKey
118 , foldlWithKey
119 -- ** Strict folds
120 , foldr'
121 , foldl'
122 , foldrWithKey'
123 , foldlWithKey'
124 -- ** Legacy folds
125 , fold
126 , foldWithKey
127
128 -- * Conversion
129 , elems
130 , keys
131 , keysSet
132 , assocs
133
134 -- ** Lists
135 , toList
136 , fromList
137 , fromListWith
138 , fromListWithKey
139
140 -- ** Ordered lists
141 , toAscList
142 , fromAscList
143 , fromAscListWith
144 , fromAscListWithKey
145 , fromDistinctAscList
146
147 -- * Filter
148 , filter
149 , filterWithKey
150 , partition
151 , partitionWithKey
152
153 , mapMaybe
154 , mapMaybeWithKey
155 , mapEither
156 , mapEitherWithKey
157
158 , split
159 , splitLookup
160
161 -- * Submap
162 , isSubmapOf, isSubmapOfBy
163 , isProperSubmapOf, isProperSubmapOfBy
164
165 -- * Min\/Max
166 , findMin
167 , findMax
168 , deleteMin
169 , deleteMax
170 , deleteFindMin
171 , deleteFindMax
172 , updateMin
173 , updateMax
174 , updateMinWithKey
175 , updateMaxWithKey
176 , minView
177 , maxView
178 , minViewWithKey
179 , maxViewWithKey
180
181 -- * Debugging
182 , showTree
183 , showTreeWith
184 ) where
185
186 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
187 import Data.Bits
188 import qualified Data.IntSet as IntSet
189 import Data.Monoid (Monoid(..))
190 import Data.Maybe (fromMaybe)
191 import Data.Typeable
192 import Data.Foldable (Foldable(foldMap))
193 import Data.Traversable (Traversable(traverse))
194 import Control.Applicative (Applicative(pure,(<*>)),(<$>))
195 import Control.Monad ( liftM )
196 {-
197 -- just for testing
198 import qualified Prelude
199 import Test.QuickCheck
200 import List (nub,sort)
201 import qualified List
202 -}
203
204 #if __GLASGOW_HASKELL__
205 import Text.Read
206 import Data.Data (Data(..), mkNoRepType)
207 #endif
208
209 #if __GLASGOW_HASKELL__ >= 503
210 import GHC.Exts ( Word(..), Int(..), shiftRL# )
211 #elif __GLASGOW_HASKELL__
212 import Word
213 import GlaExts ( Word(..), Int(..), shiftRL# )
214 #else
215 import Data.Word
216 #endif
217
218 -- Use macros to define strictness of functions.
219 -- STRICT_x_OF_y denotes an y-ary function strict in the x-th parameter.
220 -- We do not use BangPatterns, because they are not in any standard and we
221 -- want the compilers to be compiled by as many compilers as possible.
222 #define STRICT_1_OF_2(fn) fn arg _ | arg `seq` False = undefined
223
224 infixl 9 \\{-This comment teaches CPP correct behaviour -}
225
226 -- A "Nat" is a natural machine word (an unsigned Int)
227 type Nat = Word
228
229 natFromInt :: Key -> Nat
230 natFromInt = fromIntegral
231 {-# INLINE natFromInt #-}
232
233 intFromNat :: Nat -> Key
234 intFromNat = fromIntegral
235 {-# INLINE intFromNat #-}
236
237 shiftRL :: Nat -> Key -> Nat
238 #if __GLASGOW_HASKELL__
239 {--------------------------------------------------------------------
240 GHC: use unboxing to get @shiftRL@ inlined.
241 --------------------------------------------------------------------}
242 shiftRL (W# x) (I# i)
243 = W# (shiftRL# x i)
244 #else
245 shiftRL x i = shiftR x i
246 {-# INLINE shiftRL #-}
247 #endif
248
249 {--------------------------------------------------------------------
250 Operators
251 --------------------------------------------------------------------}
252
253 -- | /O(min(n,W))/. Find the value at a key.
254 -- Calls 'error' when the element can not be found.
255 --
256 -- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map
257 -- > fromList [(5,'a'), (3,'b')] ! 5 == 'a'
258
259 (!) :: IntMap a -> Key -> a
260 m ! k = find k m
261
262 -- | Same as 'difference'.
263 (\\) :: IntMap a -> IntMap b -> IntMap a
264 m1 \\ m2 = difference m1 m2
265
266 {--------------------------------------------------------------------
267 Types
268 --------------------------------------------------------------------}
269
270 -- The order of constructors of IntMap matters when considering performance.
271 -- Currently in GHC 7.0, when type has 3 constructors, they are matched from
272 -- the first to the last -- the best performance is achieved when the
273 -- constructors are ordered by frequency.
274 -- On GHC 7.0, reordering constructors from Nil | Tip | Bin to Bin | Tip | Nil
275 -- improves the containers_benchmark by 9.5% on x86 and by 8% on x86_64.
276
277 -- | A map of integers to values @a@.
278 data IntMap a = Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(IntMap a) !(IntMap a)
279 | Tip {-# UNPACK #-} !Key a
280 | Nil
281
282 type Prefix = Int
283 type Mask = Int
284 type Key = Int
285
286 instance Monoid (IntMap a) where
287 mempty = empty
288 mappend = union
289 mconcat = unions
290
291 instance Foldable IntMap where
292 foldMap _ Nil = mempty
293 foldMap f (Tip _k v) = f v
294 foldMap f (Bin _ _ l r) = foldMap f l `mappend` foldMap f r
295
296 instance Traversable IntMap where
297 traverse _ Nil = pure Nil
298 traverse f (Tip k v) = Tip k <$> f v
299 traverse f (Bin p m l r) = Bin p m <$> traverse f l <*> traverse f r
300
301 #if __GLASGOW_HASKELL__
302
303 {--------------------------------------------------------------------
304 A Data instance
305 --------------------------------------------------------------------}
306
307 -- This instance preserves data abstraction at the cost of inefficiency.
308 -- We omit reflection services for the sake of data abstraction.
309
310 instance Data a => Data (IntMap a) where
311 gfoldl f z im = z fromList `f` (toList im)
312 toConstr _ = error "toConstr"
313 gunfold _ _ = error "gunfold"
314 dataTypeOf _ = mkNoRepType "Data.IntMap.IntMap"
315 dataCast1 f = gcast1 f
316
317 #endif
318
319 {--------------------------------------------------------------------
320 Query
321 --------------------------------------------------------------------}
322 -- | /O(1)/. Is the map empty?
323 --
324 -- > Data.IntMap.null (empty) == True
325 -- > Data.IntMap.null (singleton 1 'a') == False
326
327 null :: IntMap a -> Bool
328 null Nil = True
329 null _ = False
330
331 -- | /O(n)/. Number of elements in the map.
332 --
333 -- > size empty == 0
334 -- > size (singleton 1 'a') == 1
335 -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3
336 size :: IntMap a -> Int
337 size t
338 = case t of
339 Bin _ _ l r -> size l + size r
340 Tip _ _ -> 1
341 Nil -> 0
342
343 -- | /O(min(n,W))/. Is the key a member of the map?
344 --
345 -- > member 5 (fromList [(5,'a'), (3,'b')]) == True
346 -- > member 1 (fromList [(5,'a'), (3,'b')]) == False
347
348 member :: Key -> IntMap a -> Bool
349 member k m
350 = case lookup k m of
351 Nothing -> False
352 Just _ -> True
353
354 -- | /O(log n)/. Is the key not a member of the map?
355 --
356 -- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False
357 -- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True
358
359 notMember :: Key -> IntMap a -> Bool
360 notMember k m = not $ member k m
361
362 -- The 'go' function in the lookup causes 10% speedup, but also an increased
363 -- memory allocation. It does not cause speedup with other methods like insert
364 -- and delete, so it is present only in lookup.
365
366 -- | /O(min(n,W))/. Lookup the value at a key in the map. See also 'Data.Map.lookup'.
367 lookup :: Key -> IntMap a -> Maybe a
368 lookup k = k `seq` go
369 where
370 go (Bin _ m l r)
371 | zero k m = go l
372 | otherwise = go r
373 go (Tip kx x)
374 | k == kx = Just x
375 | otherwise = Nothing
376 go Nil = Nothing
377
378
379 find :: Key -> IntMap a -> a
380 find k m
381 = case lookup k m of
382 Nothing -> error ("IntMap.find: key " ++ show k ++ " is not an element of the map")
383 Just x -> x
384
385 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
386 -- returns the value at key @k@ or returns @def@ when the key is not an
387 -- element of the map.
388 --
389 -- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x'
390 -- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a'
391
392 findWithDefault :: a -> Key -> IntMap a -> a
393 findWithDefault def k m
394 = case lookup k m of
395 Nothing -> def
396 Just x -> x
397
398 {--------------------------------------------------------------------
399 Construction
400 --------------------------------------------------------------------}
401 -- | /O(1)/. The empty map.
402 --
403 -- > empty == fromList []
404 -- > size empty == 0
405
406 empty :: IntMap a
407 empty
408 = Nil
409
410 -- | /O(1)/. A map of one element.
411 --
412 -- > singleton 1 'a' == fromList [(1, 'a')]
413 -- > size (singleton 1 'a') == 1
414
415 singleton :: Key -> a -> IntMap a
416 singleton k x
417 = Tip k x
418
419 {--------------------------------------------------------------------
420 Insert
421 --------------------------------------------------------------------}
422 -- | /O(min(n,W))/. Insert a new key\/value pair in the map.
423 -- If the key is already present in the map, the associated value is
424 -- replaced with the supplied value, i.e. 'insert' is equivalent to
425 -- @'insertWith' 'const'@.
426 --
427 -- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')]
428 -- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')]
429 -- > insert 5 'x' empty == singleton 5 'x'
430
431 insert :: Key -> a -> IntMap a -> IntMap a
432 insert k x t = k `seq`
433 case t of
434 Bin p m l r
435 | nomatch k p m -> join k (Tip k x) p t
436 | zero k m -> Bin p m (insert k x l) r
437 | otherwise -> Bin p m l (insert k x r)
438 Tip ky _
439 | k==ky -> Tip k x
440 | otherwise -> join k (Tip k x) ky t
441 Nil -> Tip k x
442
443 -- right-biased insertion, used by 'union'
444 -- | /O(min(n,W))/. Insert with a combining function.
445 -- @'insertWith' f key value mp@
446 -- will insert the pair (key, value) into @mp@ if key does
447 -- not exist in the map. If the key does exist, the function will
448 -- insert @f new_value old_value@.
449 --
450 -- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")]
451 -- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
452 -- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx"
453
454 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
455 insertWith f k x t
456 = insertWithKey (\_ x' y' -> f x' y') k x t
457
458 -- | Same as 'insertWith', but the combining function is applied strictly.
459 insertWith' :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
460 insertWith' f k x t
461 = insertWithKey' (\_ x' y' -> f x' y') k x t
462
463 -- | /O(min(n,W))/. Insert with a combining function.
464 -- @'insertWithKey' f key value mp@
465 -- will insert the pair (key, value) into @mp@ if key does
466 -- not exist in the map. If the key does exist, the function will
467 -- insert @f key new_value old_value@.
468 --
469 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
470 -- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")]
471 -- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
472 -- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx"
473
474 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
475 insertWithKey f k x t = k `seq`
476 case t of
477 Bin p m l r
478 | nomatch k p m -> join k (Tip k x) p t
479 | zero k m -> Bin p m (insertWithKey f k x l) r
480 | otherwise -> Bin p m l (insertWithKey f k x r)
481 Tip ky y
482 | k==ky -> Tip k (f k x y)
483 | otherwise -> join k (Tip k x) ky t
484 Nil -> Tip k x
485
486 -- | Same as 'insertWithKey', but the combining function is applied strictly.
487 insertWithKey' :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
488 insertWithKey' f k x t = k `seq`
489 case t of
490 Bin p m l r
491 | nomatch k p m -> join k (Tip k x) p t
492 | zero k m -> Bin p m (insertWithKey' f k x l) r
493 | otherwise -> Bin p m l (insertWithKey' f k x r)
494 Tip ky y
495 | k==ky -> let x' = f k x y in seq x' (Tip k x')
496 | otherwise -> join k (Tip k x) ky t
497 Nil -> Tip k x
498
499 -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
500 -- is a pair where the first element is equal to (@'lookup' k map@)
501 -- and the second element equal to (@'insertWithKey' f k x map@).
502 --
503 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
504 -- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")])
505 -- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")])
506 -- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx")
507 --
508 -- This is how to define @insertLookup@ using @insertLookupWithKey@:
509 --
510 -- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t
511 -- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")])
512 -- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")])
513
514 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
515 insertLookupWithKey f k x t = k `seq`
516 case t of
517 Bin p m l r
518 | nomatch k p m -> (Nothing,join k (Tip k x) p t)
519 | zero k m -> let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r)
520 | otherwise -> let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r')
521 Tip ky y
522 | k==ky -> (Just y,Tip k (f k x y))
523 | otherwise -> (Nothing,join k (Tip k x) ky t)
524 Nil -> (Nothing,Tip k x)
525
526
527 {--------------------------------------------------------------------
528 Deletion
529 [delete] is the inlined version of [deleteWith (\k x -> Nothing)]
530 --------------------------------------------------------------------}
531 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
532 -- a member of the map, the original map is returned.
533 --
534 -- > delete 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
535 -- > delete 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
536 -- > delete 5 empty == empty
537
538 delete :: Key -> IntMap a -> IntMap a
539 delete k t = k `seq`
540 case t of
541 Bin p m l r
542 | nomatch k p m -> t
543 | zero k m -> bin p m (delete k l) r
544 | otherwise -> bin p m l (delete k r)
545 Tip ky _
546 | k==ky -> Nil
547 | otherwise -> t
548 Nil -> Nil
549
550 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
551 -- a member of the map, the original map is returned.
552 --
553 -- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
554 -- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
555 -- > adjust ("new " ++) 7 empty == empty
556
557 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
558 adjust f k m
559 = adjustWithKey (\_ x -> f x) k m
560
561 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
562 -- a member of the map, the original map is returned.
563 --
564 -- > let f key x = (show key) ++ ":new " ++ x
565 -- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
566 -- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
567 -- > adjustWithKey f 7 empty == empty
568
569 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
570 adjustWithKey f
571 = updateWithKey (\k' x -> Just (f k' x))
572
573 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
574 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
575 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
576 --
577 -- > let f x = if x == "a" then Just "new a" else Nothing
578 -- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
579 -- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
580 -- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
581
582 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
583 update f
584 = updateWithKey (\_ x -> f x)
585
586 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
587 -- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
588 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
589 --
590 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
591 -- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
592 -- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
593 -- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
594
595 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
596 updateWithKey f k t = k `seq`
597 case t of
598 Bin p m l r
599 | nomatch k p m -> t
600 | zero k m -> bin p m (updateWithKey f k l) r
601 | otherwise -> bin p m l (updateWithKey f k r)
602 Tip ky y
603 | k==ky -> case (f k y) of
604 Just y' -> Tip ky y'
605 Nothing -> Nil
606 | otherwise -> t
607 Nil -> Nil
608
609 -- | /O(min(n,W))/. Lookup and update.
610 -- The function returns original value, if it is updated.
611 -- This is different behavior than 'Data.Map.updateLookupWithKey'.
612 -- Returns the original key value if the map entry is deleted.
613 --
614 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
615 -- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:new a")])
616 -- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")])
617 -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a")
618
619 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
620 updateLookupWithKey f k t = k `seq`
621 case t of
622 Bin p m l r
623 | nomatch k p m -> (Nothing,t)
624 | zero k m -> let (found,l') = updateLookupWithKey f k l in (found,bin p m l' r)
625 | otherwise -> let (found,r') = updateLookupWithKey f k r in (found,bin p m l r')
626 Tip ky y
627 | k==ky -> case (f k y) of
628 Just y' -> (Just y,Tip ky y')
629 Nothing -> (Just y,Nil)
630 | otherwise -> (Nothing,t)
631 Nil -> (Nothing,Nil)
632
633
634
635 -- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
636 -- 'alter' can be used to insert, delete, or update a value in an 'IntMap'.
637 -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
638 alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
639 alter f k t = k `seq`
640 case t of
641 Bin p m l r
642 | nomatch k p m -> case f Nothing of
643 Nothing -> t
644 Just x -> join k (Tip k x) p t
645 | zero k m -> bin p m (alter f k l) r
646 | otherwise -> bin p m l (alter f k r)
647 Tip ky y
648 | k==ky -> case f (Just y) of
649 Just x -> Tip ky x
650 Nothing -> Nil
651 | otherwise -> case f Nothing of
652 Just x -> join k (Tip k x) ky t
653 Nothing -> Tip ky y
654 Nil -> case f Nothing of
655 Just x -> Tip k x
656 Nothing -> Nil
657
658
659 {--------------------------------------------------------------------
660 Union
661 --------------------------------------------------------------------}
662 -- | The union of a list of maps.
663 --
664 -- > unions [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
665 -- > == fromList [(3, "b"), (5, "a"), (7, "C")]
666 -- > unions [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])]
667 -- > == fromList [(3, "B3"), (5, "A3"), (7, "C")]
668
669 unions :: [IntMap a] -> IntMap a
670 unions xs
671 = foldlStrict union empty xs
672
673 -- | The union of a list of maps, with a combining operation.
674 --
675 -- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
676 -- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")]
677
678 unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a
679 unionsWith f ts
680 = foldlStrict (unionWith f) empty ts
681
682 -- | /O(n+m)/. The (left-biased) union of two maps.
683 -- It prefers the first map when duplicate keys are encountered,
684 -- i.e. (@'union' == 'unionWith' 'const'@).
685 --
686 -- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")]
687
688 union :: IntMap a -> IntMap a -> IntMap a
689 union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
690 | shorter m1 m2 = union1
691 | shorter m2 m1 = union2
692 | p1 == p2 = Bin p1 m1 (union l1 l2) (union r1 r2)
693 | otherwise = join p1 t1 p2 t2
694 where
695 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
696 | zero p2 m1 = Bin p1 m1 (union l1 t2) r1
697 | otherwise = Bin p1 m1 l1 (union r1 t2)
698
699 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
700 | zero p1 m2 = Bin p2 m2 (union t1 l2) r2
701 | otherwise = Bin p2 m2 l2 (union t1 r2)
702
703 union (Tip k x) t = insert k x t
704 union t (Tip k x) = insertWith (\_ y -> y) k x t -- right bias
705 union Nil t = t
706 union t Nil = t
707
708 -- | /O(n+m)/. The union with a combining function.
709 --
710 -- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")]
711
712 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
713 unionWith f m1 m2
714 = unionWithKey (\_ x y -> f x y) m1 m2
715
716 -- | /O(n+m)/. The union with a combining function.
717 --
718 -- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value
719 -- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")]
720
721 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
722 unionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
723 | shorter m1 m2 = union1
724 | shorter m2 m1 = union2
725 | p1 == p2 = Bin p1 m1 (unionWithKey f l1 l2) (unionWithKey f r1 r2)
726 | otherwise = join p1 t1 p2 t2
727 where
728 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
729 | zero p2 m1 = Bin p1 m1 (unionWithKey f l1 t2) r1
730 | otherwise = Bin p1 m1 l1 (unionWithKey f r1 t2)
731
732 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
733 | zero p1 m2 = Bin p2 m2 (unionWithKey f t1 l2) r2
734 | otherwise = Bin p2 m2 l2 (unionWithKey f t1 r2)
735
736 unionWithKey f (Tip k x) t = insertWithKey f k x t
737 unionWithKey f t (Tip k x) = insertWithKey (\k' x' y' -> f k' y' x') k x t -- right bias
738 unionWithKey _ Nil t = t
739 unionWithKey _ t Nil = t
740
741 {--------------------------------------------------------------------
742 Difference
743 --------------------------------------------------------------------}
744 -- | /O(n+m)/. Difference between two maps (based on keys).
745 --
746 -- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b"
747
748 difference :: IntMap a -> IntMap b -> IntMap a
749 difference t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
750 | shorter m1 m2 = difference1
751 | shorter m2 m1 = difference2
752 | p1 == p2 = bin p1 m1 (difference l1 l2) (difference r1 r2)
753 | otherwise = t1
754 where
755 difference1 | nomatch p2 p1 m1 = t1
756 | zero p2 m1 = bin p1 m1 (difference l1 t2) r1
757 | otherwise = bin p1 m1 l1 (difference r1 t2)
758
759 difference2 | nomatch p1 p2 m2 = t1
760 | zero p1 m2 = difference t1 l2
761 | otherwise = difference t1 r2
762
763 difference t1@(Tip k _) t2
764 | member k t2 = Nil
765 | otherwise = t1
766
767 difference Nil _ = Nil
768 difference t (Tip k _) = delete k t
769 difference t Nil = t
770
771 -- | /O(n+m)/. Difference with a combining function.
772 --
773 -- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing
774 -- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")])
775 -- > == singleton 3 "b:B"
776
777 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
778 differenceWith f m1 m2
779 = differenceWithKey (\_ x y -> f x y) m1 m2
780
781 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
782 -- encountered, the combining function is applied to the key and both values.
783 -- If it returns 'Nothing', the element is discarded (proper set difference).
784 -- If it returns (@'Just' y@), the element is updated with a new value @y@.
785 --
786 -- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing
787 -- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")])
788 -- > == singleton 3 "3:b|B"
789
790 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
791 differenceWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
792 | shorter m1 m2 = difference1
793 | shorter m2 m1 = difference2
794 | p1 == p2 = bin p1 m1 (differenceWithKey f l1 l2) (differenceWithKey f r1 r2)
795 | otherwise = t1
796 where
797 difference1 | nomatch p2 p1 m1 = t1
798 | zero p2 m1 = bin p1 m1 (differenceWithKey f l1 t2) r1
799 | otherwise = bin p1 m1 l1 (differenceWithKey f r1 t2)
800
801 difference2 | nomatch p1 p2 m2 = t1
802 | zero p1 m2 = differenceWithKey f t1 l2
803 | otherwise = differenceWithKey f t1 r2
804
805 differenceWithKey f t1@(Tip k x) t2
806 = case lookup k t2 of
807 Just y -> case f k x y of
808 Just y' -> Tip k y'
809 Nothing -> Nil
810 Nothing -> t1
811
812 differenceWithKey _ Nil _ = Nil
813 differenceWithKey f t (Tip k y) = updateWithKey (\k' x -> f k' x y) k t
814 differenceWithKey _ t Nil = t
815
816
817 {--------------------------------------------------------------------
818 Intersection
819 --------------------------------------------------------------------}
820 -- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys).
821 --
822 -- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a"
823
824 intersection :: IntMap a -> IntMap b -> IntMap a
825 intersection t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
826 | shorter m1 m2 = intersection1
827 | shorter m2 m1 = intersection2
828 | p1 == p2 = bin p1 m1 (intersection l1 l2) (intersection r1 r2)
829 | otherwise = Nil
830 where
831 intersection1 | nomatch p2 p1 m1 = Nil
832 | zero p2 m1 = intersection l1 t2
833 | otherwise = intersection r1 t2
834
835 intersection2 | nomatch p1 p2 m2 = Nil
836 | zero p1 m2 = intersection t1 l2
837 | otherwise = intersection t1 r2
838
839 intersection t1@(Tip k _) t2
840 | member k t2 = t1
841 | otherwise = Nil
842 intersection t (Tip k _)
843 = case lookup k t of
844 Just y -> Tip k y
845 Nothing -> Nil
846 intersection Nil _ = Nil
847 intersection _ Nil = Nil
848
849 -- | /O(n+m)/. The intersection with a combining function.
850 --
851 -- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA"
852
853 intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
854 intersectionWith f m1 m2
855 = intersectionWithKey (\_ x y -> f x y) m1 m2
856
857 -- | /O(n+m)/. The intersection with a combining function.
858 --
859 -- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar
860 -- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A"
861
862 intersectionWithKey :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
863 intersectionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
864 | shorter m1 m2 = intersection1
865 | shorter m2 m1 = intersection2
866 | p1 == p2 = bin p1 m1 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2)
867 | otherwise = Nil
868 where
869 intersection1 | nomatch p2 p1 m1 = Nil
870 | zero p2 m1 = intersectionWithKey f l1 t2
871 | otherwise = intersectionWithKey f r1 t2
872
873 intersection2 | nomatch p1 p2 m2 = Nil
874 | zero p1 m2 = intersectionWithKey f t1 l2
875 | otherwise = intersectionWithKey f t1 r2
876
877 intersectionWithKey f (Tip k x) t2
878 = case lookup k t2 of
879 Just y -> Tip k (f k x y)
880 Nothing -> Nil
881 intersectionWithKey f t1 (Tip k y)
882 = case lookup k t1 of
883 Just x -> Tip k (f k x y)
884 Nothing -> Nil
885 intersectionWithKey _ Nil _ = Nil
886 intersectionWithKey _ _ Nil = Nil
887
888
889 {--------------------------------------------------------------------
890 Min\/Max
891 --------------------------------------------------------------------}
892
893 -- | /O(log n)/. Update the value at the minimal key.
894 --
895 -- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"3:b"), (5,"a")]
896 -- > updateMinWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
897
898 updateMinWithKey :: (Key -> a -> a) -> IntMap a -> IntMap a
899 updateMinWithKey f t
900 = case t of
901 Bin p m l r | m < 0 -> let t' = updateMinWithKeyUnsigned f r in Bin p m l t'
902 Bin p m l r -> let t' = updateMinWithKeyUnsigned f l in Bin p m t' r
903 Tip k y -> Tip k (f k y)
904 Nil -> error "maxView: empty map has no maximal element"
905
906 updateMinWithKeyUnsigned :: (Key -> a -> a) -> IntMap a -> IntMap a
907 updateMinWithKeyUnsigned f t
908 = case t of
909 Bin p m l r -> let t' = updateMinWithKeyUnsigned f l in Bin p m t' r
910 Tip k y -> Tip k (f k y)
911 Nil -> error "updateMinWithKeyUnsigned Nil"
912
913 -- | /O(log n)/. Update the value at the maximal key.
914 --
915 -- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")]
916 -- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
917
918 updateMaxWithKey :: (Key -> a -> a) -> IntMap a -> IntMap a
919 updateMaxWithKey f t
920 = case t of
921 Bin p m l r | m < 0 -> let t' = updateMaxWithKeyUnsigned f l in Bin p m t' r
922 Bin p m l r -> let t' = updateMaxWithKeyUnsigned f r in Bin p m l t'
923 Tip k y -> Tip k (f k y)
924 Nil -> error "maxView: empty map has no maximal element"
925
926 updateMaxWithKeyUnsigned :: (Key -> a -> a) -> IntMap a -> IntMap a
927 updateMaxWithKeyUnsigned f t
928 = case t of
929 Bin p m l r -> let t' = updateMaxWithKeyUnsigned f r in Bin p m l t'
930 Tip k y -> Tip k (f k y)
931 Nil -> error "updateMaxWithKeyUnsigned Nil"
932
933
934 -- | /O(log n)/. Retrieves the maximal (key,value) pair of the map, and
935 -- the map stripped of that element, or 'Nothing' if passed an empty map.
936 --
937 -- > maxViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((5,"a"), singleton 3 "b")
938 -- > maxViewWithKey empty == Nothing
939
940 maxViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
941 maxViewWithKey t
942 = case t of
943 Bin p m l r | m < 0 -> let (result, t') = maxViewUnsigned l in Just (result, bin p m t' r)
944 Bin p m l r -> let (result, t') = maxViewUnsigned r in Just (result, bin p m l t')
945 Tip k y -> Just ((k,y), Nil)
946 Nil -> Nothing
947
948 maxViewUnsigned :: IntMap a -> ((Key, a), IntMap a)
949 maxViewUnsigned t
950 = case t of
951 Bin p m l r -> let (result,t') = maxViewUnsigned r in (result,bin p m l t')
952 Tip k y -> ((k,y), Nil)
953 Nil -> error "maxViewUnsigned Nil"
954
955 -- | /O(log n)/. Retrieves the minimal (key,value) pair of the map, and
956 -- the map stripped of that element, or 'Nothing' if passed an empty map.
957 --
958 -- > minViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((3,"b"), singleton 5 "a")
959 -- > minViewWithKey empty == Nothing
960
961 minViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
962 minViewWithKey t
963 = case t of
964 Bin p m l r | m < 0 -> let (result, t') = minViewUnsigned r in Just (result, bin p m l t')
965 Bin p m l r -> let (result, t') = minViewUnsigned l in Just (result, bin p m t' r)
966 Tip k y -> Just ((k,y),Nil)
967 Nil -> Nothing
968
969 minViewUnsigned :: IntMap a -> ((Key, a), IntMap a)
970 minViewUnsigned t
971 = case t of
972 Bin p m l r -> let (result,t') = minViewUnsigned l in (result,bin p m t' r)
973 Tip k y -> ((k,y),Nil)
974 Nil -> error "minViewUnsigned Nil"
975
976
977 -- | /O(log n)/. Update the value at the maximal key.
978 --
979 -- > updateMax (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "Xa")]
980 -- > updateMax (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
981
982 updateMax :: (a -> a) -> IntMap a -> IntMap a
983 updateMax f = updateMaxWithKey (const f)
984
985 -- | /O(log n)/. Update the value at the minimal key.
986 --
987 -- > updateMin (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "Xb"), (5, "a")]
988 -- > updateMin (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
989
990 updateMin :: (a -> a) -> IntMap a -> IntMap a
991 updateMin f = updateMinWithKey (const f)
992
993 -- Similar to the Arrow instance.
994 first :: (a -> c) -> (a, b) -> (c, b)
995 first f (x,y) = (f x,y)
996
997 -- | /O(log n)/. Retrieves the maximal key of the map, and the map
998 -- stripped of that element, or 'Nothing' if passed an empty map.
999 maxView :: IntMap a -> Maybe (a, IntMap a)
1000 maxView t = liftM (first snd) (maxViewWithKey t)
1001
1002 -- | /O(log n)/. Retrieves the minimal key of the map, and the map
1003 -- stripped of that element, or 'Nothing' if passed an empty map.
1004 minView :: IntMap a -> Maybe (a, IntMap a)
1005 minView t = liftM (first snd) (minViewWithKey t)
1006
1007 -- | /O(log n)/. Delete and find the maximal element.
1008 deleteFindMax :: IntMap a -> (a, IntMap a)
1009 deleteFindMax = fromMaybe (error "deleteFindMax: empty map has no maximal element") . maxView
1010
1011 -- | /O(log n)/. Delete and find the minimal element.
1012 deleteFindMin :: IntMap a -> (a, IntMap a)
1013 deleteFindMin = fromMaybe (error "deleteFindMin: empty map has no minimal element") . minView
1014
1015 -- | /O(log n)/. The minimal key of the map.
1016 findMin :: IntMap a -> (Key, a)
1017 findMin Nil = error $ "findMin: empty map has no minimal element"
1018 findMin (Tip k v) = (k,v)
1019 findMin (Bin _ m l r)
1020 | m < 0 = go r
1021 | otherwise = go l
1022 where go (Tip k v) = (k,v)
1023 go (Bin _ _ l' _) = go l'
1024 go Nil = error "findMax Nil"
1025
1026 -- | /O(log n)/. The maximal key of the map.
1027 findMax :: IntMap a -> (Key, a)
1028 findMax Nil = error $ "findMax: empty map has no maximal element"
1029 findMax (Tip k v) = (k,v)
1030 findMax (Bin _ m l r)
1031 | m < 0 = go l
1032 | otherwise = go r
1033 where go (Tip k v) = (k,v)
1034 go (Bin _ _ _ r') = go r'
1035 go Nil = error "findMax Nil"
1036
1037 -- | /O(log n)/. Delete the minimal key. An error is thrown if the IntMap is already empty.
1038 -- Note, this is not the same behavior Map.
1039 deleteMin :: IntMap a -> IntMap a
1040 deleteMin = maybe (error "deleteMin: empty map has no minimal element") snd . minView
1041
1042 -- | /O(log n)/. Delete the maximal key. An error is thrown if the IntMap is already empty.
1043 -- Note, this is not the same behavior Map.
1044 deleteMax :: IntMap a -> IntMap a
1045 deleteMax = maybe (error "deleteMax: empty map has no maximal element") snd . maxView
1046
1047
1048 {--------------------------------------------------------------------
1049 Submap
1050 --------------------------------------------------------------------}
1051 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
1052 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
1053 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
1054 isProperSubmapOf m1 m2
1055 = isProperSubmapOfBy (==) m1 m2
1056
1057 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
1058 The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
1059 @m1@ and @m2@ are not equal,
1060 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
1061 applied to their respective values. For example, the following
1062 expressions are all 'True':
1063
1064 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1065 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1066
1067 But the following are all 'False':
1068
1069 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
1070 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
1071 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1072 -}
1073 isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
1074 isProperSubmapOfBy predicate t1 t2
1075 = case submapCmp predicate t1 t2 of
1076 LT -> True
1077 _ -> False
1078
1079 submapCmp :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Ordering
1080 submapCmp predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1081 | shorter m1 m2 = GT
1082 | shorter m2 m1 = submapCmpLt
1083 | p1 == p2 = submapCmpEq
1084 | otherwise = GT -- disjoint
1085 where
1086 submapCmpLt | nomatch p1 p2 m2 = GT
1087 | zero p1 m2 = submapCmp predicate t1 l2
1088 | otherwise = submapCmp predicate t1 r2
1089 submapCmpEq = case (submapCmp predicate l1 l2, submapCmp predicate r1 r2) of
1090 (GT,_ ) -> GT
1091 (_ ,GT) -> GT
1092 (EQ,EQ) -> EQ
1093 _ -> LT
1094
1095 submapCmp _ (Bin _ _ _ _) _ = GT
1096 submapCmp predicate (Tip kx x) (Tip ky y)
1097 | (kx == ky) && predicate x y = EQ
1098 | otherwise = GT -- disjoint
1099 submapCmp predicate (Tip k x) t
1100 = case lookup k t of
1101 Just y | predicate x y -> LT
1102 _ -> GT -- disjoint
1103 submapCmp _ Nil Nil = EQ
1104 submapCmp _ Nil _ = LT
1105
1106 -- | /O(n+m)/. Is this a submap?
1107 -- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
1108 isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
1109 isSubmapOf m1 m2
1110 = isSubmapOfBy (==) m1 m2
1111
1112 {- | /O(n+m)/.
1113 The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if
1114 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
1115 applied to their respective values. For example, the following
1116 expressions are all 'True':
1117
1118 > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1119 > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1120 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
1121
1122 But the following are all 'False':
1123
1124 > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)])
1125 > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1126 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
1127 -}
1128 isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
1129 isSubmapOfBy predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1130 | shorter m1 m2 = False
1131 | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy predicate t1 l2
1132 else isSubmapOfBy predicate t1 r2)
1133 | otherwise = (p1==p2) && isSubmapOfBy predicate l1 l2 && isSubmapOfBy predicate r1 r2
1134 isSubmapOfBy _ (Bin _ _ _ _) _ = False
1135 isSubmapOfBy predicate (Tip k x) t = case lookup k t of
1136 Just y -> predicate x y
1137 Nothing -> False
1138 isSubmapOfBy _ Nil _ = True
1139
1140 {--------------------------------------------------------------------
1141 Mapping
1142 --------------------------------------------------------------------}
1143 -- | /O(n)/. Map a function over all values in the map.
1144 --
1145 -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")]
1146
1147 map :: (a -> b) -> IntMap a -> IntMap b
1148 map f = mapWithKey (\_ x -> f x)
1149
1150 -- | /O(n)/. Map a function over all values in the map.
1151 --
1152 -- > let f key x = (show key) ++ ":" ++ x
1153 -- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")]
1154
1155 mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b
1156 mapWithKey f t
1157 = case t of
1158 Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r)
1159 Tip k x -> Tip k (f k x)
1160 Nil -> Nil
1161
1162 -- | /O(n)/. The function @'mapAccum'@ threads an accumulating
1163 -- argument through the map in ascending order of keys.
1164 --
1165 -- > let f a b = (a ++ b, b ++ "X")
1166 -- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")])
1167
1168 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1169 mapAccum f = mapAccumWithKey (\a' _ x -> f a' x)
1170
1171 -- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating
1172 -- argument through the map in ascending order of keys.
1173 --
1174 -- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X")
1175 -- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")])
1176
1177 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1178 mapAccumWithKey f a t
1179 = mapAccumL f a t
1180
1181 -- | /O(n)/. The function @'mapAccumL'@ threads an accumulating
1182 -- argument through the map in ascending order of keys.
1183 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1184 mapAccumL f a t
1185 = case t of
1186 Bin p m l r -> let (a1,l') = mapAccumL f a l
1187 (a2,r') = mapAccumL f a1 r
1188 in (a2,Bin p m l' r')
1189 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
1190 Nil -> (a,Nil)
1191
1192 -- | /O(n)/. The function @'mapAccumR'@ threads an accumulating
1193 -- argument through the map in descending order of keys.
1194 mapAccumRWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1195 mapAccumRWithKey f a t
1196 = case t of
1197 Bin p m l r -> let (a1,r') = mapAccumRWithKey f a r
1198 (a2,l') = mapAccumRWithKey f a1 l
1199 in (a2,Bin p m l' r')
1200 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
1201 Nil -> (a,Nil)
1202
1203 {--------------------------------------------------------------------
1204 Filter
1205 --------------------------------------------------------------------}
1206 -- | /O(n)/. Filter all values that satisfy some predicate.
1207 --
1208 -- > filter (> "a") (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1209 -- > filter (> "x") (fromList [(5,"a"), (3,"b")]) == empty
1210 -- > filter (< "a") (fromList [(5,"a"), (3,"b")]) == empty
1211
1212 filter :: (a -> Bool) -> IntMap a -> IntMap a
1213 filter p m
1214 = filterWithKey (\_ x -> p x) m
1215
1216 -- | /O(n)/. Filter all keys\/values that satisfy some predicate.
1217 --
1218 -- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1219
1220 filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a
1221 filterWithKey predicate t
1222 = case t of
1223 Bin p m l r
1224 -> bin p m (filterWithKey predicate l) (filterWithKey predicate r)
1225 Tip k x
1226 | predicate k x -> t
1227 | otherwise -> Nil
1228 Nil -> Nil
1229
1230 -- | /O(n)/. Partition the map according to some predicate. The first
1231 -- map contains all elements that satisfy the predicate, the second all
1232 -- elements that fail the predicate. See also 'split'.
1233 --
1234 -- > partition (> "a") (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
1235 -- > partition (< "x") (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
1236 -- > partition (> "x") (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
1237
1238 partition :: (a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
1239 partition p m
1240 = partitionWithKey (\_ x -> p x) m
1241
1242 -- | /O(n)/. Partition the map according to some predicate. The first
1243 -- map contains all elements that satisfy the predicate, the second all
1244 -- elements that fail the predicate. See also 'split'.
1245 --
1246 -- > partitionWithKey (\ k _ -> k > 3) (fromList [(5,"a"), (3,"b")]) == (singleton 5 "a", singleton 3 "b")
1247 -- > partitionWithKey (\ k _ -> k < 7) (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
1248 -- > partitionWithKey (\ k _ -> k > 7) (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
1249
1250 partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
1251 partitionWithKey predicate t
1252 = case t of
1253 Bin p m l r
1254 -> let (l1,l2) = partitionWithKey predicate l
1255 (r1,r2) = partitionWithKey predicate r
1256 in (bin p m l1 r1, bin p m l2 r2)
1257 Tip k x
1258 | predicate k x -> (t,Nil)
1259 | otherwise -> (Nil,t)
1260 Nil -> (Nil,Nil)
1261
1262 -- | /O(n)/. Map values and collect the 'Just' results.
1263 --
1264 -- > let f x = if x == "a" then Just "new a" else Nothing
1265 -- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a"
1266
1267 mapMaybe :: (a -> Maybe b) -> IntMap a -> IntMap b
1268 mapMaybe f = mapMaybeWithKey (\_ x -> f x)
1269
1270 -- | /O(n)/. Map keys\/values and collect the 'Just' results.
1271 --
1272 -- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing
1273 -- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3"
1274
1275 mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b
1276 mapMaybeWithKey f (Bin p m l r)
1277 = bin p m (mapMaybeWithKey f l) (mapMaybeWithKey f r)
1278 mapMaybeWithKey f (Tip k x) = case f k x of
1279 Just y -> Tip k y
1280 Nothing -> Nil
1281 mapMaybeWithKey _ Nil = Nil
1282
1283 -- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
1284 --
1285 -- > let f a = if a < "c" then Left a else Right a
1286 -- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1287 -- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")])
1288 -- >
1289 -- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1290 -- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1291
1292 mapEither :: (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
1293 mapEither f m
1294 = mapEitherWithKey (\_ x -> f x) m
1295
1296 -- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
1297 --
1298 -- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a)
1299 -- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1300 -- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")])
1301 -- >
1302 -- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1303 -- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")])
1304
1305 mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
1306 mapEitherWithKey f (Bin p m l r)
1307 = (bin p m l1 r1, bin p m l2 r2)
1308 where
1309 (l1,l2) = mapEitherWithKey f l
1310 (r1,r2) = mapEitherWithKey f r
1311 mapEitherWithKey f (Tip k x) = case f k x of
1312 Left y -> (Tip k y, Nil)
1313 Right z -> (Nil, Tip k z)
1314 mapEitherWithKey _ Nil = (Nil, Nil)
1315
1316 -- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@
1317 -- where all keys in @map1@ are lower than @k@ and all keys in
1318 -- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
1319 --
1320 -- > split 2 (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3,"b"), (5,"a")])
1321 -- > split 3 (fromList [(5,"a"), (3,"b")]) == (empty, singleton 5 "a")
1322 -- > split 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
1323 -- > split 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", empty)
1324 -- > split 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], empty)
1325
1326 split :: Key -> IntMap a -> (IntMap a,IntMap a)
1327 split k t
1328 = case t of
1329 Bin _ m l r
1330 | m < 0 -> (if k >= 0 -- handle negative numbers.
1331 then let (lt,gt) = split' k l in (union r lt, gt)
1332 else let (lt,gt) = split' k r in (lt, union gt l))
1333 | otherwise -> split' k t
1334 Tip ky _
1335 | k>ky -> (t,Nil)
1336 | k<ky -> (Nil,t)
1337 | otherwise -> (Nil,Nil)
1338 Nil -> (Nil,Nil)
1339
1340 split' :: Key -> IntMap a -> (IntMap a,IntMap a)
1341 split' k t
1342 = case t of
1343 Bin p m l r
1344 | nomatch k p m -> if k>p then (t,Nil) else (Nil,t)
1345 | zero k m -> let (lt,gt) = split k l in (lt,union gt r)
1346 | otherwise -> let (lt,gt) = split k r in (union l lt,gt)
1347 Tip ky _
1348 | k>ky -> (t,Nil)
1349 | k<ky -> (Nil,t)
1350 | otherwise -> (Nil,Nil)
1351 Nil -> (Nil,Nil)
1352
1353 -- | /O(log n)/. Performs a 'split' but also returns whether the pivot
1354 -- key was found in the original map.
1355 --
1356 -- > splitLookup 2 (fromList [(5,"a"), (3,"b")]) == (empty, Nothing, fromList [(3,"b"), (5,"a")])
1357 -- > splitLookup 3 (fromList [(5,"a"), (3,"b")]) == (empty, Just "b", singleton 5 "a")
1358 -- > splitLookup 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Nothing, singleton 5 "a")
1359 -- > splitLookup 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Just "a", empty)
1360 -- > splitLookup 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], Nothing, empty)
1361
1362 splitLookup :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
1363 splitLookup k t
1364 = case t of
1365 Bin _ m l r
1366 | m < 0 -> (if k >= 0 -- handle negative numbers.
1367 then let (lt,found,gt) = splitLookup' k l in (union r lt,found, gt)
1368 else let (lt,found,gt) = splitLookup' k r in (lt,found, union gt l))
1369 | otherwise -> splitLookup' k t
1370 Tip ky y
1371 | k>ky -> (t,Nothing,Nil)
1372 | k<ky -> (Nil,Nothing,t)
1373 | otherwise -> (Nil,Just y,Nil)
1374 Nil -> (Nil,Nothing,Nil)
1375
1376 splitLookup' :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
1377 splitLookup' k t
1378 = case t of
1379 Bin p m l r
1380 | nomatch k p m -> if k>p then (t,Nothing,Nil) else (Nil,Nothing,t)
1381 | zero k m -> let (lt,found,gt) = splitLookup k l in (lt,found,union gt r)
1382 | otherwise -> let (lt,found,gt) = splitLookup k r in (union l lt,found,gt)
1383 Tip ky y
1384 | k>ky -> (t,Nothing,Nil)
1385 | k<ky -> (Nil,Nothing,t)
1386 | otherwise -> (Nil,Just y,Nil)
1387 Nil -> (Nil,Nothing,Nil)
1388
1389 {--------------------------------------------------------------------
1390 Fold
1391 --------------------------------------------------------------------}
1392 -- | /O(n)/. Fold the values in the map using the given right-associative
1393 -- binary operator. This function is an equivalent of 'foldr' and is present
1394 -- for compatibility only.
1395 --
1396 -- /Please note that fold will be deprecated in the future and removed./
1397 fold :: (a -> b -> b) -> b -> IntMap a -> b
1398 fold = foldr
1399 {-# INLINE fold #-}
1400
1401 -- | /O(n)/. Fold the values in the map using the given right-associative
1402 -- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'elems'@.
1403 --
1404 -- For example,
1405 --
1406 -- > elems map = foldr (:) [] map
1407 --
1408 -- > let f a len = len + (length a)
1409 -- > foldr f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
1410 foldr :: (a -> b -> b) -> b -> IntMap a -> b
1411 foldr f z t =
1412 case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before
1413 _ -> go z t
1414 where
1415 go z' Nil = z'
1416 go z' (Tip _ x) = f x z'
1417 go z' (Bin _ _ l r) = go (go z' r) l
1418 {-# INLINE foldr #-}
1419
1420 -- | /O(n)/. A strict version of 'foldr'. Each application of the operator is
1421 -- evaluated before using the result in the next application. This
1422 -- function is strict in the starting value.
1423 foldr' :: (a -> b -> b) -> b -> IntMap a -> b
1424 foldr' f z t =
1425 case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before
1426 _ -> go z t
1427 where
1428 STRICT_1_OF_2(go)
1429 go z' Nil = z'
1430 go z' (Tip _ x) = f x z'
1431 go z' (Bin _ _ l r) = go (go z' r) l
1432 {-# INLINE foldr' #-}
1433
1434 -- | /O(n)/. Fold the values in the map using the given left-associative
1435 -- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@.
1436 --
1437 -- For example,
1438 --
1439 -- > elems = reverse . foldl (flip (:)) []
1440 --
1441 -- > let f len a = len + (length a)
1442 -- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
1443 foldl :: (a -> b -> a) -> a -> IntMap b -> a
1444 foldl f z t =
1445 case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before
1446 _ -> go z t
1447 where
1448 go z' Nil = z'
1449 go z' (Tip _ x) = f z' x
1450 go z' (Bin _ _ l r) = go (go z' l) r
1451 {-# INLINE foldl #-}
1452
1453 -- | /O(n)/. A strict version of 'foldl'. Each application of the operator is
1454 -- evaluated before using the result in the next application. This
1455 -- function is strict in the starting value.
1456 foldl' :: (a -> b -> a) -> a -> IntMap b -> a
1457 foldl' f z t =
1458 case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before
1459 _ -> go z t
1460 where
1461 STRICT_1_OF_2(go)
1462 go z' Nil = z'
1463 go z' (Tip _ x) = f z' x
1464 go z' (Bin _ _ l r) = go (go z' l) r
1465 {-# INLINE foldl' #-}
1466
1467 -- | /O(n)/. Fold the keys and values in the map using the given right-associative
1468 -- binary operator. This function is an equivalent of 'foldrWithKey' and is present
1469 -- for compatibility only.
1470 --
1471 -- /Please note that foldWithKey will be deprecated in the future and removed./
1472 foldWithKey :: (Int -> a -> b -> b) -> b -> IntMap a -> b
1473 foldWithKey = foldrWithKey
1474 {-# INLINE foldWithKey #-}
1475
1476 -- | /O(n)/. Fold the keys and values in the map using the given right-associative
1477 -- binary operator, such that
1478 -- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
1479 --
1480 -- For example,
1481 --
1482 -- > keys map = foldrWithKey (\k x ks -> k:ks) [] map
1483 --
1484 -- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
1485 -- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)"
1486 foldrWithKey :: (Int -> a -> b -> b) -> b -> IntMap a -> b
1487 foldrWithKey f z t =
1488 case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before
1489 _ -> go z t
1490 where
1491 go z' Nil = z'
1492 go z' (Tip kx x) = f kx x z'
1493 go z' (Bin _ _ l r) = go (go z' r) l
1494 {-# INLINE foldrWithKey #-}
1495
1496 -- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is
1497 -- evaluated before using the result in the next application. This
1498 -- function is strict in the starting value.
1499 foldrWithKey' :: (Int -> a -> b -> b) -> b -> IntMap a -> b
1500 foldrWithKey' f z t =
1501 case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before
1502 _ -> go z t
1503 where
1504 STRICT_1_OF_2(go)
1505 go z' Nil = z'
1506 go z' (Tip kx x) = f kx x z'
1507 go z' (Bin _ _ l r) = go (go z' r) l
1508 {-# INLINE foldrWithKey' #-}
1509
1510 -- | /O(n)/. Fold the keys and values in the map using the given left-associative
1511 -- binary operator, such that
1512 -- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@.
1513 --
1514 -- For example,
1515 --
1516 -- > keys = reverse . foldlWithKey (\ks k x -> k:ks) []
1517 --
1518 -- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
1519 -- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)"
1520 foldlWithKey :: (a -> Int -> b -> a) -> a -> IntMap b -> a
1521 foldlWithKey f z t =
1522 case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before
1523 _ -> go z t
1524 where
1525 go z' Nil = z'
1526 go z' (Tip kx x) = f z' kx x
1527 go z' (Bin _ _ l r) = go (go z' l) r
1528 {-# INLINE foldlWithKey #-}
1529
1530 -- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is
1531 -- evaluated before using the result in the next application. This
1532 -- function is strict in the starting value.
1533 foldlWithKey' :: (a -> Int -> b -> a) -> a -> IntMap b -> a
1534 foldlWithKey' f z t =
1535 case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before
1536 _ -> go z t
1537 where
1538 STRICT_1_OF_2(go)
1539 go z' Nil = z'
1540 go z' (Tip kx x) = f z' kx x
1541 go z' (Bin _ _ l r) = go (go z' l) r
1542 {-# INLINE foldlWithKey' #-}
1543
1544 {--------------------------------------------------------------------
1545 List variations
1546 --------------------------------------------------------------------}
1547 -- | /O(n)/.
1548 -- Return all elements of the map in the ascending order of their keys.
1549 --
1550 -- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"]
1551 -- > elems empty == []
1552
1553 elems :: IntMap a -> [a]
1554 elems
1555 = foldr (:) []
1556
1557 -- | /O(n)/. Return all keys of the map in ascending order.
1558 --
1559 -- > keys (fromList [(5,"a"), (3,"b")]) == [3,5]
1560 -- > keys empty == []
1561
1562 keys :: IntMap a -> [Key]
1563 keys
1564 = foldrWithKey (\k _ ks -> k:ks) []
1565
1566 -- | /O(n*min(n,W))/. The set of all keys of the map.
1567 --
1568 -- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.IntSet.fromList [3,5]
1569 -- > keysSet empty == Data.IntSet.empty
1570
1571 keysSet :: IntMap a -> IntSet.IntSet
1572 keysSet m = IntSet.fromDistinctAscList (keys m)
1573
1574
1575 -- | /O(n)/. Return all key\/value pairs in the map in ascending key order.
1576 --
1577 -- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1578 -- > assocs empty == []
1579
1580 assocs :: IntMap a -> [(Key,a)]
1581 assocs m
1582 = toList m
1583
1584
1585 {--------------------------------------------------------------------
1586 Lists
1587 --------------------------------------------------------------------}
1588 -- | /O(n)/. Convert the map to a list of key\/value pairs.
1589 --
1590 -- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1591 -- > toList empty == []
1592
1593 toList :: IntMap a -> [(Key,a)]
1594 toList
1595 = foldrWithKey (\k x xs -> (k,x):xs) []
1596
1597 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
1598 -- keys are in ascending order.
1599 --
1600 -- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1601
1602 toAscList :: IntMap a -> [(Key,a)]
1603 toAscList t
1604 = -- NOTE: the following algorithm only works for big-endian trees
1605 let (pos,neg) = span (\(k,_) -> k >=0) (foldrWithKey (\k x xs -> (k,x):xs) [] t) in neg ++ pos
1606
1607 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
1608 --
1609 -- > fromList [] == empty
1610 -- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
1611 -- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
1612
1613 fromList :: [(Key,a)] -> IntMap a
1614 fromList xs
1615 = foldlStrict ins empty xs
1616 where
1617 ins t (k,x) = insert k x t
1618
1619 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
1620 --
1621 -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")]
1622 -- > fromListWith (++) [] == empty
1623
1624 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1625 fromListWith f xs
1626 = fromListWithKey (\_ x y -> f x y) xs
1627
1628 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
1629 --
1630 -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")]
1631 -- > fromListWith (++) [] == empty
1632
1633 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1634 fromListWithKey f xs
1635 = foldlStrict ins empty xs
1636 where
1637 ins t (k,x) = insertWithKey f k x t
1638
1639 -- | /O(n)/. Build a map from a list of key\/value pairs where
1640 -- the keys are in ascending order.
1641 --
1642 -- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
1643 -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")]
1644
1645 fromAscList :: [(Key,a)] -> IntMap a
1646 fromAscList xs
1647 = fromAscListWithKey (\_ x _ -> x) xs
1648
1649 -- | /O(n)/. Build a map from a list of key\/value pairs where
1650 -- the keys are in ascending order, with a combining function on equal keys.
1651 -- /The precondition (input list is ascending) is not checked./
1652 --
1653 -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
1654
1655 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1656 fromAscListWith f xs
1657 = fromAscListWithKey (\_ x y -> f x y) xs
1658
1659 -- | /O(n)/. Build a map from a list of key\/value pairs where
1660 -- the keys are in ascending order, with a combining function on equal keys.
1661 -- /The precondition (input list is ascending) is not checked./
1662 --
1663 -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
1664
1665 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1666 fromAscListWithKey _ [] = Nil
1667 fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0)
1668 where
1669 -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
1670 combineEq z [] = [z]
1671 combineEq z@(kz,zz) (x@(kx,xx):xs)
1672 | kx==kz = let yy = f kx xx zz in combineEq (kx,yy) xs
1673 | otherwise = z:combineEq x xs
1674
1675 -- | /O(n)/. Build a map from a list of key\/value pairs where
1676 -- the keys are in ascending order and all distinct.
1677 -- /The precondition (input list is strictly ascending) is not checked./
1678 --
1679 -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
1680
1681 #ifdef __GLASGOW_HASKELL__
1682 fromDistinctAscList :: forall a. [(Key,a)] -> IntMap a
1683 #else
1684 fromDistinctAscList :: [(Key,a)] -> IntMap a
1685 #endif
1686 fromDistinctAscList [] = Nil
1687 fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada
1688 where
1689 work (kx,vx) [] stk = finish kx (Tip kx vx) stk
1690 work (kx,vx) (z@(kz,_):zs) stk = reduce z zs (branchMask kx kz) kx (Tip kx vx) stk
1691
1692 #ifdef __GLASGOW_HASKELL__
1693 reduce :: (Key,a) -> [(Key,a)] -> Mask -> Prefix -> IntMap a -> Stack a -> IntMap a
1694 #endif
1695 reduce z zs _ px tx Nada = work z zs (Push px tx Nada)
1696 reduce z zs m px tx stk@(Push py ty stk') =
1697 let mxy = branchMask px py
1698 pxy = mask px mxy
1699 in if shorter m mxy
1700 then reduce z zs m pxy (Bin pxy mxy ty tx) stk'
1701 else work z zs (Push px tx stk)
1702
1703 finish _ t Nada = t
1704 finish px tx (Push py ty stk) = finish p (join py ty px tx) stk
1705 where m = branchMask px py
1706 p = mask px m
1707
1708 data Stack a = Push {-# UNPACK #-} !Prefix !(IntMap a) !(Stack a) | Nada
1709
1710
1711 {--------------------------------------------------------------------
1712 Eq
1713 --------------------------------------------------------------------}
1714 instance Eq a => Eq (IntMap a) where
1715 t1 == t2 = equal t1 t2
1716 t1 /= t2 = nequal t1 t2
1717
1718 equal :: Eq a => IntMap a -> IntMap a -> Bool
1719 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1720 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
1721 equal (Tip kx x) (Tip ky y)
1722 = (kx == ky) && (x==y)
1723 equal Nil Nil = True
1724 equal _ _ = False
1725
1726 nequal :: Eq a => IntMap a -> IntMap a -> Bool
1727 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1728 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
1729 nequal (Tip kx x) (Tip ky y)
1730 = (kx /= ky) || (x/=y)
1731 nequal Nil Nil = False
1732 nequal _ _ = True
1733
1734 {--------------------------------------------------------------------
1735 Ord
1736 --------------------------------------------------------------------}
1737
1738 instance Ord a => Ord (IntMap a) where
1739 compare m1 m2 = compare (toList m1) (toList m2)
1740
1741 {--------------------------------------------------------------------
1742 Functor
1743 --------------------------------------------------------------------}
1744
1745 instance Functor IntMap where
1746 fmap = map
1747
1748 {--------------------------------------------------------------------
1749 Show
1750 --------------------------------------------------------------------}
1751
1752 instance Show a => Show (IntMap a) where
1753 showsPrec d m = showParen (d > 10) $
1754 showString "fromList " . shows (toList m)
1755
1756 {-
1757 XXX unused code
1758
1759 showMap :: (Show a) => [(Key,a)] -> ShowS
1760 showMap []
1761 = showString "{}"
1762 showMap (x:xs)
1763 = showChar '{' . showElem x . showTail xs
1764 where
1765 showTail [] = showChar '}'
1766 showTail (x':xs') = showChar ',' . showElem x' . showTail xs'
1767
1768 showElem (k,v) = shows k . showString ":=" . shows v
1769 -}
1770
1771 {--------------------------------------------------------------------
1772 Read
1773 --------------------------------------------------------------------}
1774 instance (Read e) => Read (IntMap e) where
1775 #ifdef __GLASGOW_HASKELL__
1776 readPrec = parens $ prec 10 $ do
1777 Ident "fromList" <- lexP
1778 xs <- readPrec
1779 return (fromList xs)
1780
1781 readListPrec = readListPrecDefault
1782 #else
1783 readsPrec p = readParen (p > 10) $ \ r -> do
1784 ("fromList",s) <- lex r
1785 (xs,t) <- reads s
1786 return (fromList xs,t)
1787 #endif
1788
1789 {--------------------------------------------------------------------
1790 Typeable
1791 --------------------------------------------------------------------}
1792
1793 #include "Typeable.h"
1794 INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap")
1795
1796 {--------------------------------------------------------------------
1797 Debugging
1798 --------------------------------------------------------------------}
1799 -- | /O(n)/. Show the tree that implements the map. The tree is shown
1800 -- in a compressed, hanging format.
1801 showTree :: Show a => IntMap a -> String
1802 showTree s
1803 = showTreeWith True False s
1804
1805
1806 {- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows
1807 the tree that implements the map. If @hang@ is
1808 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
1809 @wide@ is 'True', an extra wide version is shown.
1810 -}
1811 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
1812 showTreeWith hang wide t
1813 | hang = (showsTreeHang wide [] t) ""
1814 | otherwise = (showsTree wide [] [] t) ""
1815
1816 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
1817 showsTree wide lbars rbars t
1818 = case t of
1819 Bin p m l r
1820 -> showsTree wide (withBar rbars) (withEmpty rbars) r .
1821 showWide wide rbars .
1822 showsBars lbars . showString (showBin p m) . showString "\n" .
1823 showWide wide lbars .
1824 showsTree wide (withEmpty lbars) (withBar lbars) l
1825 Tip k x
1826 -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1827 Nil -> showsBars lbars . showString "|\n"
1828
1829 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
1830 showsTreeHang wide bars t
1831 = case t of
1832 Bin p m l r
1833 -> showsBars bars . showString (showBin p m) . showString "\n" .
1834 showWide wide bars .
1835 showsTreeHang wide (withBar bars) l .
1836 showWide wide bars .
1837 showsTreeHang wide (withEmpty bars) r
1838 Tip k x
1839 -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1840 Nil -> showsBars bars . showString "|\n"
1841
1842 showBin :: Prefix -> Mask -> String
1843 showBin _ _
1844 = "*" -- ++ show (p,m)
1845
1846 showWide :: Bool -> [String] -> String -> String
1847 showWide wide bars
1848 | wide = showString (concat (reverse bars)) . showString "|\n"
1849 | otherwise = id
1850
1851 showsBars :: [String] -> ShowS
1852 showsBars bars
1853 = case bars of
1854 [] -> id
1855 _ -> showString (concat (reverse (tail bars))) . showString node
1856
1857 node :: String
1858 node = "+--"
1859
1860 withBar, withEmpty :: [String] -> [String]
1861 withBar bars = "| ":bars
1862 withEmpty bars = " ":bars
1863
1864
1865 {--------------------------------------------------------------------
1866 Helpers
1867 --------------------------------------------------------------------}
1868 {--------------------------------------------------------------------
1869 Join
1870 --------------------------------------------------------------------}
1871 join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
1872 join p1 t1 p2 t2
1873 | zero p1 m = Bin p m t1 t2
1874 | otherwise = Bin p m t2 t1
1875 where
1876 m = branchMask p1 p2
1877 p = mask p1 m
1878 {-# INLINE join #-}
1879
1880 {--------------------------------------------------------------------
1881 @bin@ assures that we never have empty trees within a tree.
1882 --------------------------------------------------------------------}
1883 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
1884 bin _ _ l Nil = l
1885 bin _ _ Nil r = r
1886 bin p m l r = Bin p m l r
1887 {-# INLINE bin #-}
1888
1889
1890 {--------------------------------------------------------------------
1891 Endian independent bit twiddling
1892 --------------------------------------------------------------------}
1893 zero :: Key -> Mask -> Bool
1894 zero i m
1895 = (natFromInt i) .&. (natFromInt m) == 0
1896 {-# INLINE zero #-}
1897
1898 nomatch,match :: Key -> Prefix -> Mask -> Bool
1899 nomatch i p m
1900 = (mask i m) /= p
1901 {-# INLINE nomatch #-}
1902
1903 match i p m
1904 = (mask i m) == p
1905 {-# INLINE match #-}
1906
1907 mask :: Key -> Mask -> Prefix
1908 mask i m
1909 = maskW (natFromInt i) (natFromInt m)
1910 {-# INLINE mask #-}
1911
1912
1913 {--------------------------------------------------------------------
1914 Big endian operations
1915 --------------------------------------------------------------------}
1916 maskW :: Nat -> Nat -> Prefix
1917 maskW i m
1918 = intFromNat (i .&. (complement (m-1) `xor` m))
1919 {-# INLINE maskW #-}
1920
1921 shorter :: Mask -> Mask -> Bool
1922 shorter m1 m2
1923 = (natFromInt m1) > (natFromInt m2)
1924 {-# INLINE shorter #-}
1925
1926 branchMask :: Prefix -> Prefix -> Mask
1927 branchMask p1 p2
1928 = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
1929 {-# INLINE branchMask #-}
1930
1931 {----------------------------------------------------------------------
1932 Finding the highest bit (mask) in a word [x] can be done efficiently in
1933 three ways:
1934 * convert to a floating point value and the mantissa tells us the
1935 [log2(x)] that corresponds with the highest bit position. The mantissa
1936 is retrieved either via the standard C function [frexp] or by some bit
1937 twiddling on IEEE compatible numbers (float). Note that one needs to
1938 use at least [double] precision for an accurate mantissa of 32 bit
1939 numbers.
1940 * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
1941 * use processor specific assembler instruction (asm).
1942
1943 The most portable way would be [bit], but is it efficient enough?
1944 I have measured the cycle counts of the different methods on an AMD
1945 Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
1946
1947 highestBitMask: method cycles
1948 --------------
1949 frexp 200
1950 float 33
1951 bit 11
1952 asm 12
1953
1954 highestBit: method cycles
1955 --------------
1956 frexp 195
1957 float 33
1958 bit 11
1959 asm 11
1960
1961 Wow, the bit twiddling is on today's RISC like machines even faster
1962 than a single CISC instruction (BSR)!
1963 ----------------------------------------------------------------------}
1964
1965 {----------------------------------------------------------------------
1966 [highestBitMask] returns a word where only the highest bit is set.
1967 It is found by first setting all bits in lower positions than the
1968 highest bit and than taking an exclusive or with the original value.
1969 Allthough the function may look expensive, GHC compiles this into
1970 excellent C code that subsequently compiled into highly efficient
1971 machine code. The algorithm is derived from Jorg Arndt's FXT library.
1972 ----------------------------------------------------------------------}
1973 highestBitMask :: Nat -> Nat
1974 highestBitMask x0
1975 = case (x0 .|. shiftRL x0 1) of
1976 x1 -> case (x1 .|. shiftRL x1 2) of
1977 x2 -> case (x2 .|. shiftRL x2 4) of
1978 x3 -> case (x3 .|. shiftRL x3 8) of
1979 x4 -> case (x4 .|. shiftRL x4 16) of
1980 x5 -> case (x5 .|. shiftRL x5 32) of -- for 64 bit platforms
1981 x6 -> (x6 `xor` (shiftRL x6 1))
1982 {-# INLINE highestBitMask #-}
1983
1984
1985 {--------------------------------------------------------------------
1986 Utilities
1987 --------------------------------------------------------------------}
1988
1989 foldlStrict :: (a -> b -> a) -> a -> [b] -> a
1990 foldlStrict f = go
1991 where
1992 go z [] = z
1993 go z (x:xs) = let z' = f z x in z' `seq` go z' xs
1994 {-# INLINE foldlStrict #-}