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