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