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