Add restrictKeys and withoutKeys
[packages/containers.git] / Data / IntMap / Base.hs
1 {-# LANGUAGE CPP #-}
2 {-# LANGUAGE BangPatterns #-}
3 #if __GLASGOW_HASKELL__
4 {-# LANGUAGE MagicHash, DeriveDataTypeable, StandaloneDeriving #-}
5 {-# LANGUAGE ScopedTypeVariables #-}
6 #endif
7 #if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703
8 {-# LANGUAGE Trustworthy #-}
9 #endif
10 #if __GLASGOW_HASKELL__ >= 708
11 {-# LANGUAGE TypeFamilies #-}
12 #endif
13
14 #include "containers.h"
15
16 -----------------------------------------------------------------------------
17 -- |
18 -- Module : Data.IntMap.Base
19 -- Copyright : (c) Daan Leijen 2002
20 -- (c) Andriy Palamarchuk 2008
21 -- License : BSD-style
22 -- Maintainer : libraries@haskell.org
23 -- Stability : provisional
24 -- Portability : portable
25 --
26 -- This defines the data structures and core (hidden) manipulations
27 -- on representations.
28 -----------------------------------------------------------------------------
29
30 -- [Note: INLINE bit fiddling]
31 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~
32 -- It is essential that the bit fiddling functions like mask, zero, branchMask
33 -- etc are inlined. If they do not, the memory allocation skyrockets. The GHC
34 -- usually gets it right, but it is disastrous if it does not. Therefore we
35 -- explicitly mark these functions INLINE.
36
37
38 -- [Note: Local 'go' functions and capturing]
39 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
40 -- Care must be taken when using 'go' function which captures an argument.
41 -- Sometimes (for example when the argument is passed to a data constructor,
42 -- as in insert), GHC heap-allocates more than necessary. Therefore C-- code
43 -- must be checked for increased allocation when creating and modifying such
44 -- functions.
45
46
47 -- [Note: Order of constructors]
48 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
49 -- The order of constructors of IntMap matters when considering performance.
50 -- Currently in GHC 7.0, when type has 3 constructors, they are matched from
51 -- the first to the last -- the best performance is achieved when the
52 -- constructors are ordered by frequency.
53 -- On GHC 7.0, reordering constructors from Nil | Tip | Bin to Bin | Tip | Nil
54 -- improves the benchmark by circa 10%.
55
56 module Data.IntMap.Base (
57 -- * Map type
58 IntMap(..), Key -- instance Eq,Show
59
60 -- * Operators
61 , (!), (\\)
62
63 -- * Query
64 , null
65 , size
66 , member
67 , notMember
68 , lookup
69 , findWithDefault
70 , lookupLT
71 , lookupGT
72 , lookupLE
73 , lookupGE
74
75 -- * Construction
76 , empty
77 , singleton
78
79 -- ** Insertion
80 , insert
81 , insertWith
82 , insertWithKey
83 , insertLookupWithKey
84
85 -- ** Delete\/Update
86 , delete
87 , adjust
88 , adjustWithKey
89 , update
90 , updateWithKey
91 , updateLookupWithKey
92 , alter
93 , alterF
94
95 -- * Combine
96
97 -- ** Union
98 , union
99 , unionWith
100 , unionWithKey
101 , unions
102 , unionsWith
103
104 -- ** Difference
105 , difference
106 , differenceWith
107 , differenceWithKey
108
109 -- ** Intersection
110 , intersection
111 , intersectionWith
112 , intersectionWithKey
113
114 -- ** Universal combining function
115 , mergeWithKey
116 , mergeWithKey'
117
118 -- * Traversal
119 -- ** Map
120 , map
121 , mapWithKey
122 , traverseWithKey
123 , mapAccum
124 , mapAccumWithKey
125 , mapAccumRWithKey
126 , mapKeys
127 , mapKeysWith
128 , mapKeysMonotonic
129
130 -- * Folds
131 , foldr
132 , foldl
133 , foldrWithKey
134 , foldlWithKey
135 , foldMapWithKey
136
137 -- ** Strict folds
138 , foldr'
139 , foldl'
140 , foldrWithKey'
141 , foldlWithKey'
142
143 -- * Conversion
144 , elems
145 , keys
146 , assocs
147 , keysSet
148 , fromSet
149
150 -- ** Lists
151 , toList
152 , fromList
153 , fromListWith
154 , fromListWithKey
155
156 -- ** Ordered lists
157 , toAscList
158 , toDescList
159 , fromAscList
160 , fromAscListWith
161 , fromAscListWithKey
162 , fromDistinctAscList
163
164 -- * Filter
165 , filter
166 , filterWithKey
167 , restrictKeys
168 , withoutKeys
169 , partition
170 , partitionWithKey
171
172 , mapMaybe
173 , mapMaybeWithKey
174 , mapEither
175 , mapEitherWithKey
176
177 , split
178 , splitLookup
179 , splitRoot
180
181 -- * Submap
182 , isSubmapOf, isSubmapOfBy
183 , isProperSubmapOf, isProperSubmapOfBy
184
185 -- * Min\/Max
186 , findMin
187 , findMax
188 , deleteMin
189 , deleteMax
190 , deleteFindMin
191 , deleteFindMax
192 , updateMin
193 , updateMax
194 , updateMinWithKey
195 , updateMaxWithKey
196 , minView
197 , maxView
198 , minViewWithKey
199 , maxViewWithKey
200
201 -- * Debugging
202 , showTree
203 , showTreeWith
204
205 -- * Internal types
206 , Mask, Prefix, Nat
207
208 -- * Utility
209 , natFromInt
210 , intFromNat
211 , link
212 , bin
213 , binCheckLeft
214 , binCheckRight
215 , zero
216 , nomatch
217 , match
218 , mask
219 , maskW
220 , shorter
221 , branchMask
222 , highestBitMask
223 ) where
224
225 #if !(MIN_VERSION_base(4,8,0))
226 import Control.Applicative (Applicative(pure, (<*>)), (<$>))
227 import Data.Monoid (Monoid(..))
228 import Data.Traversable (Traversable(traverse))
229 import Data.Word (Word)
230 #endif
231 #if MIN_VERSION_base(4,9,0)
232 import Data.Semigroup (Semigroup((<>), stimes), stimesIdempotentMonoid)
233 #endif
234
235 import Control.DeepSeq (NFData(rnf))
236 import Control.Monad (liftM)
237 import Data.Bits
238 import qualified Data.Foldable as Foldable
239 import Data.Maybe (fromMaybe)
240 import Data.Typeable
241 import Prelude hiding (lookup, map, filter, foldr, foldl, null)
242
243 import Data.IntSet.Base (Key)
244 import qualified Data.IntSet.Base as IntSet
245 import Data.Utils.BitUtil
246 import Data.Utils.StrictFold
247 import Data.Utils.StrictPair
248
249 #if __GLASGOW_HASKELL__
250 import Data.Data (Data(..), Constr, mkConstr, constrIndex, Fixity(Prefix),
251 DataType, mkDataType)
252 import GHC.Exts (build)
253 import Data.Functor ((<$))
254 #if __GLASGOW_HASKELL__ >= 708
255 import qualified GHC.Exts as GHCExts
256 #endif
257 import Text.Read
258 #endif
259 #if __GLASGOW_HASKELL__ >= 709
260 import Data.Coerce
261 #endif
262
263
264 -- A "Nat" is a natural machine word (an unsigned Int)
265 type Nat = Word
266
267 natFromInt :: Key -> Nat
268 natFromInt = fromIntegral
269 {-# INLINE natFromInt #-}
270
271 intFromNat :: Nat -> Key
272 intFromNat = fromIntegral
273 {-# INLINE intFromNat #-}
274
275 {--------------------------------------------------------------------
276 Types
277 --------------------------------------------------------------------}
278
279
280 -- | A map of integers to values @a@.
281
282 -- See Note: Order of constructors
283 data IntMap a = Bin {-# UNPACK #-} !Prefix
284 {-# UNPACK #-} !Mask
285 !(IntMap a)
286 !(IntMap a)
287 | Tip {-# UNPACK #-} !Key a
288 | Nil
289
290 type Prefix = Int
291 type Mask = Int
292
293 {--------------------------------------------------------------------
294 Operators
295 --------------------------------------------------------------------}
296
297 -- | /O(min(n,W))/. Find the value at a key.
298 -- Calls 'error' when the element can not be found.
299 --
300 -- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map
301 -- > fromList [(5,'a'), (3,'b')] ! 5 == 'a'
302
303 (!) :: IntMap a -> Key -> a
304 (!) m k = find k m
305
306 -- | Same as 'difference'.
307 (\\) :: IntMap a -> IntMap b -> IntMap a
308 m1 \\ m2 = difference m1 m2
309
310 infixl 9 \\{-This comment teaches CPP correct behaviour -}
311
312 {--------------------------------------------------------------------
313 Types
314 --------------------------------------------------------------------}
315
316 instance Monoid (IntMap a) where
317 mempty = empty
318 mconcat = unions
319 #if !(MIN_VERSION_base(4,9,0))
320 mappend = union
321 #else
322 mappend = (<>)
323
324 instance Semigroup (IntMap a) where
325 (<>) = union
326 stimes = stimesIdempotentMonoid
327 #endif
328
329 instance Foldable.Foldable IntMap where
330 fold = go
331 where go Nil = mempty
332 go (Tip _ v) = v
333 go (Bin _ _ l r) = go l `mappend` go r
334 {-# INLINABLE fold #-}
335 foldr = foldr
336 {-# INLINE foldr #-}
337 foldl = foldl
338 {-# INLINE foldl #-}
339 foldMap f t = go t
340 where go Nil = mempty
341 go (Tip _ v) = f v
342 go (Bin _ _ l r) = go l `mappend` go r
343 {-# INLINE foldMap #-}
344
345 #if MIN_VERSION_base(4,6,0)
346 foldl' = foldl'
347 {-# INLINE foldl' #-}
348 foldr' = foldr'
349 {-# INLINE foldr' #-}
350 #endif
351 #if MIN_VERSION_base(4,8,0)
352 length = size
353 {-# INLINE length #-}
354 null = null
355 {-# INLINE null #-}
356 toList = elems -- NB: Foldable.toList /= IntMap.toList
357 {-# INLINE toList #-}
358 elem = go
359 where go !_ Nil = False
360 go x (Tip _ y) = x == y
361 go x (Bin _ _ l r) = go x l || go x r
362 {-# INLINABLE elem #-}
363 maximum = start
364 where start Nil = error "IntMap.Foldable.maximum: called with empty map"
365 start (Tip _ y) = y
366 start (Bin _ _ l r) = go (start l) r
367
368 go !m Nil = m
369 go m (Tip _ y) = max m y
370 go m (Bin _ _ l r) = go (go m l) r
371 {-# INLINABLE maximum #-}
372 minimum = start
373 where start Nil = error "IntMap.Foldable.minimum: called with empty map"
374 start (Tip _ y) = y
375 start (Bin _ _ l r) = go (start l) r
376
377 go !m Nil = m
378 go m (Tip _ y) = min m y
379 go m (Bin _ _ l r) = go (go m l) r
380 {-# INLINABLE minimum #-}
381 sum = foldl' (+) 0
382 {-# INLINABLE sum #-}
383 product = foldl' (*) 1
384 {-# INLINABLE product #-}
385 #endif
386
387 instance Traversable IntMap where
388 traverse f = traverseWithKey (\_ -> f)
389 {-# INLINE traverse #-}
390
391 instance NFData a => NFData (IntMap a) where
392 rnf Nil = ()
393 rnf (Tip _ v) = rnf v
394 rnf (Bin _ _ l r) = rnf l `seq` rnf r
395
396 #if __GLASGOW_HASKELL__
397
398 {--------------------------------------------------------------------
399 A Data instance
400 --------------------------------------------------------------------}
401
402 -- This instance preserves data abstraction at the cost of inefficiency.
403 -- We provide limited reflection services for the sake of data abstraction.
404
405 instance Data a => Data (IntMap a) where
406 gfoldl f z im = z fromList `f` (toList im)
407 toConstr _ = fromListConstr
408 gunfold k z c = case constrIndex c of
409 1 -> k (z fromList)
410 _ -> error "gunfold"
411 dataTypeOf _ = intMapDataType
412 dataCast1 f = gcast1 f
413
414 fromListConstr :: Constr
415 fromListConstr = mkConstr intMapDataType "fromList" [] Prefix
416
417 intMapDataType :: DataType
418 intMapDataType = mkDataType "Data.IntMap.Base.IntMap" [fromListConstr]
419
420 #endif
421
422 {--------------------------------------------------------------------
423 Query
424 --------------------------------------------------------------------}
425 -- | /O(1)/. Is the map empty?
426 --
427 -- > Data.IntMap.null (empty) == True
428 -- > Data.IntMap.null (singleton 1 'a') == False
429
430 null :: IntMap a -> Bool
431 null Nil = True
432 null _ = False
433 {-# INLINE null #-}
434
435 -- | /O(n)/. Number of elements in the map.
436 --
437 -- > size empty == 0
438 -- > size (singleton 1 'a') == 1
439 -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3
440 size :: IntMap a -> Int
441 size (Bin _ _ l r) = size l + size r
442 size (Tip _ _) = 1
443 size Nil = 0
444
445 -- | /O(min(n,W))/. Is the key a member of the map?
446 --
447 -- > member 5 (fromList [(5,'a'), (3,'b')]) == True
448 -- > member 1 (fromList [(5,'a'), (3,'b')]) == False
449
450 -- See Note: Local 'go' functions and capturing]
451 member :: Key -> IntMap a -> Bool
452 member !k = go
453 where
454 go (Bin p m l r) | nomatch k p m = False
455 | zero k m = go l
456 | otherwise = go r
457 go (Tip kx _) = k == kx
458 go Nil = False
459
460 -- | /O(min(n,W))/. Is the key not a member of the map?
461 --
462 -- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False
463 -- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True
464
465 notMember :: Key -> IntMap a -> Bool
466 notMember k m = not $ member k m
467
468 -- | /O(min(n,W))/. Lookup the value at a key in the map. See also 'Data.Map.lookup'.
469
470 -- See Note: Local 'go' functions and capturing]
471 lookup :: Key -> IntMap a -> Maybe a
472 lookup !k = go
473 where
474 go (Bin p m l r) | nomatch k p m = Nothing
475 | zero k m = go l
476 | otherwise = go r
477 go (Tip kx x) | k == kx = Just x
478 | otherwise = Nothing
479 go Nil = Nothing
480
481
482 -- See Note: Local 'go' functions and capturing]
483 find :: Key -> IntMap a -> a
484 find !k = go
485 where
486 go (Bin p m l r) | nomatch k p m = not_found
487 | zero k m = go l
488 | otherwise = go r
489 go (Tip kx x) | k == kx = x
490 | otherwise = not_found
491 go Nil = not_found
492
493 not_found = error ("IntMap.!: key " ++ show k ++ " is not an element of the map")
494
495 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
496 -- returns the value at key @k@ or returns @def@ when the key is not an
497 -- element of the map.
498 --
499 -- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x'
500 -- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a'
501
502 -- See Note: Local 'go' functions and capturing]
503 findWithDefault :: a -> Key -> IntMap a -> a
504 findWithDefault def !k = go
505 where
506 go (Bin p m l r) | nomatch k p m = def
507 | zero k m = go l
508 | otherwise = go r
509 go (Tip kx x) | k == kx = x
510 | otherwise = def
511 go Nil = def
512
513 -- | /O(log n)/. Find largest key smaller than the given one and return the
514 -- corresponding (key, value) pair.
515 --
516 -- > lookupLT 3 (fromList [(3,'a'), (5,'b')]) == Nothing
517 -- > lookupLT 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
518
519 -- See Note: Local 'go' functions and capturing.
520 lookupLT :: Key -> IntMap a -> Maybe (Key, a)
521 lookupLT !k t = case t of
522 Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r
523 _ -> go Nil t
524 where
525 go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMax def else unsafeFindMax r
526 | zero k m = go def l
527 | otherwise = go l r
528 go def (Tip ky y) | k <= ky = unsafeFindMax def
529 | otherwise = Just (ky, y)
530 go def Nil = unsafeFindMax def
531
532 -- | /O(log n)/. Find smallest key greater than the given one and return the
533 -- corresponding (key, value) pair.
534 --
535 -- > lookupGT 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
536 -- > lookupGT 5 (fromList [(3,'a'), (5,'b')]) == Nothing
537
538 -- See Note: Local 'go' functions and capturing.
539 lookupGT :: Key -> IntMap a -> Maybe (Key, a)
540 lookupGT !k t = case t of
541 Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r
542 _ -> go Nil t
543 where
544 go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMin l else unsafeFindMin def
545 | zero k m = go r l
546 | otherwise = go def r
547 go def (Tip ky y) | k >= ky = unsafeFindMin def
548 | otherwise = Just (ky, y)
549 go def Nil = unsafeFindMin def
550
551 -- | /O(log n)/. Find largest key smaller or equal to the given one and return
552 -- the corresponding (key, value) pair.
553 --
554 -- > lookupLE 2 (fromList [(3,'a'), (5,'b')]) == Nothing
555 -- > lookupLE 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
556 -- > lookupLE 5 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
557
558 -- See Note: Local 'go' functions and capturing.
559 lookupLE :: Key -> IntMap a -> Maybe (Key, a)
560 lookupLE !k t = case t of
561 Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r
562 _ -> go Nil t
563 where
564 go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMax def else unsafeFindMax r
565 | zero k m = go def l
566 | otherwise = go l r
567 go def (Tip ky y) | k < ky = unsafeFindMax def
568 | otherwise = Just (ky, y)
569 go def Nil = unsafeFindMax def
570
571 -- | /O(log n)/. Find smallest key greater or equal to the given one and return
572 -- the corresponding (key, value) pair.
573 --
574 -- > lookupGE 3 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
575 -- > lookupGE 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
576 -- > lookupGE 6 (fromList [(3,'a'), (5,'b')]) == Nothing
577
578 -- See Note: Local 'go' functions and capturing.
579 lookupGE :: Key -> IntMap a -> Maybe (Key, a)
580 lookupGE !k t = case t of
581 Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r
582 _ -> go Nil t
583 where
584 go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMin l else unsafeFindMin def
585 | zero k m = go r l
586 | otherwise = go def r
587 go def (Tip ky y) | k > ky = unsafeFindMin def
588 | otherwise = Just (ky, y)
589 go def Nil = unsafeFindMin def
590
591
592 -- Helper function for lookupGE and lookupGT. It assumes that if a Bin node is
593 -- given, it has m > 0.
594 unsafeFindMin :: IntMap a -> Maybe (Key, a)
595 unsafeFindMin Nil = Nothing
596 unsafeFindMin (Tip ky y) = Just (ky, y)
597 unsafeFindMin (Bin _ _ l _) = unsafeFindMin l
598
599 -- Helper function for lookupLE and lookupLT. It assumes that if a Bin node is
600 -- given, it has m > 0.
601 unsafeFindMax :: IntMap a -> Maybe (Key, a)
602 unsafeFindMax Nil = Nothing
603 unsafeFindMax (Tip ky y) = Just (ky, y)
604 unsafeFindMax (Bin _ _ _ r) = unsafeFindMax r
605
606 {--------------------------------------------------------------------
607 Construction
608 --------------------------------------------------------------------}
609 -- | /O(1)/. The empty map.
610 --
611 -- > empty == fromList []
612 -- > size empty == 0
613
614 empty :: IntMap a
615 empty
616 = Nil
617 {-# INLINE empty #-}
618
619 -- | /O(1)/. A map of one element.
620 --
621 -- > singleton 1 'a' == fromList [(1, 'a')]
622 -- > size (singleton 1 'a') == 1
623
624 singleton :: Key -> a -> IntMap a
625 singleton k x
626 = Tip k x
627 {-# INLINE singleton #-}
628
629 {--------------------------------------------------------------------
630 Insert
631 --------------------------------------------------------------------}
632 -- | /O(min(n,W))/. Insert a new key\/value pair in the map.
633 -- If the key is already present in the map, the associated value is
634 -- replaced with the supplied value, i.e. 'insert' is equivalent to
635 -- @'insertWith' 'const'@.
636 --
637 -- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')]
638 -- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')]
639 -- > insert 5 'x' empty == singleton 5 'x'
640
641 insert :: Key -> a -> IntMap a -> IntMap a
642 insert !k x t@(Bin p m l r)
643 | nomatch k p m = link k (Tip k x) p t
644 | zero k m = Bin p m (insert k x l) r
645 | otherwise = Bin p m l (insert k x r)
646 insert k x t@(Tip ky _)
647 | k==ky = Tip k x
648 | otherwise = link k (Tip k x) ky t
649 insert k x Nil = Tip k x
650
651 -- right-biased insertion, used by 'union'
652 -- | /O(min(n,W))/. Insert with a combining function.
653 -- @'insertWith' f key value mp@
654 -- will insert the pair (key, value) into @mp@ if key does
655 -- not exist in the map. If the key does exist, the function will
656 -- insert @f new_value old_value@.
657 --
658 -- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")]
659 -- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
660 -- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx"
661
662 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
663 insertWith f k x t
664 = insertWithKey (\_ x' y' -> f x' y') k x t
665
666 -- | /O(min(n,W))/. Insert with a combining function.
667 -- @'insertWithKey' f key value mp@
668 -- will insert the pair (key, value) into @mp@ if key does
669 -- not exist in the map. If the key does exist, the function will
670 -- insert @f key new_value old_value@.
671 --
672 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
673 -- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")]
674 -- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
675 -- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx"
676
677 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
678 insertWithKey f !k x t@(Bin p m l r)
679 | nomatch k p m = link k (Tip k x) p t
680 | zero k m = Bin p m (insertWithKey f k x l) r
681 | otherwise = Bin p m l (insertWithKey f k x r)
682 insertWithKey f k x t@(Tip ky y)
683 | k == ky = Tip k (f k x y)
684 | otherwise = link k (Tip k x) ky t
685 insertWithKey _ k x Nil = Tip k x
686
687 -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
688 -- is a pair where the first element is equal to (@'lookup' k map@)
689 -- and the second element equal to (@'insertWithKey' f k x map@).
690 --
691 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
692 -- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")])
693 -- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")])
694 -- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx")
695 --
696 -- This is how to define @insertLookup@ using @insertLookupWithKey@:
697 --
698 -- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t
699 -- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")])
700 -- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")])
701
702 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
703 insertLookupWithKey f !k x t@(Bin p m l r)
704 | nomatch k p m = (Nothing,link k (Tip k x) p t)
705 | zero k m = let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r)
706 | otherwise = let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r')
707 insertLookupWithKey f k x t@(Tip ky y)
708 | k == ky = (Just y,Tip k (f k x y))
709 | otherwise = (Nothing,link k (Tip k x) ky t)
710 insertLookupWithKey _ k x Nil = (Nothing,Tip k x)
711
712
713 {--------------------------------------------------------------------
714 Deletion
715 --------------------------------------------------------------------}
716 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
717 -- a member of the map, the original map is returned.
718 --
719 -- > delete 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
720 -- > delete 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
721 -- > delete 5 empty == empty
722
723 delete :: Key -> IntMap a -> IntMap a
724 delete !k t@(Bin p m l r)
725 | nomatch k p m = t
726 | zero k m = binCheckLeft p m (delete k l) r
727 | otherwise = binCheckRight p m l (delete k r)
728 delete k t@(Tip ky _)
729 | k == ky = Nil
730 | otherwise = t
731 delete _k Nil = Nil
732
733 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
734 -- a member of the map, the original map is returned.
735 --
736 -- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
737 -- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
738 -- > adjust ("new " ++) 7 empty == empty
739
740 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
741 adjust f k m
742 = adjustWithKey (\_ x -> f x) k m
743
744 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
745 -- a member of the map, the original map is returned.
746 --
747 -- > let f key x = (show key) ++ ":new " ++ x
748 -- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
749 -- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
750 -- > adjustWithKey f 7 empty == empty
751
752 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
753 adjustWithKey f !k t@(Bin p m l r)
754 | nomatch k p m = t
755 | zero k m = Bin p m (adjustWithKey f k l) r
756 | otherwise = Bin p m l (adjustWithKey f k r)
757 adjustWithKey f k t@(Tip ky y)
758 | k == ky = Tip ky (f k y)
759 | otherwise = t
760 adjustWithKey _ _ Nil = Nil
761
762
763 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
764 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
765 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
766 --
767 -- > let f x = if x == "a" then Just "new a" else Nothing
768 -- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
769 -- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
770 -- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
771
772 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
773 update f
774 = updateWithKey (\_ x -> f x)
775
776 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
777 -- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
778 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
779 --
780 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
781 -- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
782 -- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
783 -- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
784
785 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
786 updateWithKey f !k t@(Bin p m l r)
787 | nomatch k p m = t
788 | zero k m = binCheckLeft p m (updateWithKey f k l) r
789 | otherwise = binCheckRight p m l (updateWithKey f k r)
790 updateWithKey f k t@(Tip ky y)
791 | k == ky = case (f k y) of
792 Just y' -> Tip ky y'
793 Nothing -> Nil
794 | otherwise = t
795 updateWithKey _ _ Nil = Nil
796
797 -- | /O(min(n,W))/. Lookup and update.
798 -- The function returns original value, if it is updated.
799 -- This is different behavior than 'Data.Map.updateLookupWithKey'.
800 -- Returns the original key value if the map entry is deleted.
801 --
802 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
803 -- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:new a")])
804 -- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")])
805 -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a")
806
807 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
808 updateLookupWithKey f !k t@(Bin p m l r)
809 | nomatch k p m = (Nothing,t)
810 | zero k m = let !(found,l') = updateLookupWithKey f k l in (found,binCheckLeft p m l' r)
811 | otherwise = let !(found,r') = updateLookupWithKey f k r in (found,binCheckRight p m l r')
812 updateLookupWithKey f k t@(Tip ky y)
813 | k==ky = case (f k y) of
814 Just y' -> (Just y,Tip ky y')
815 Nothing -> (Just y,Nil)
816 | otherwise = (Nothing,t)
817 updateLookupWithKey _ _ Nil = (Nothing,Nil)
818
819
820
821 -- | /O(min(n,W))/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
822 -- 'alter' can be used to insert, delete, or update a value in an 'IntMap'.
823 -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
824 alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
825 alter f !k t@(Bin p m l r)
826 | nomatch k p m = case f Nothing of
827 Nothing -> t
828 Just x -> link k (Tip k x) p t
829 | zero k m = binCheckLeft p m (alter f k l) r
830 | otherwise = binCheckRight p m l (alter f k r)
831 alter f k t@(Tip ky y)
832 | k==ky = case f (Just y) of
833 Just x -> Tip ky x
834 Nothing -> Nil
835 | otherwise = case f Nothing of
836 Just x -> link k (Tip k x) ky t
837 Nothing -> Tip ky y
838 alter f k Nil = case f Nothing of
839 Just x -> Tip k x
840 Nothing -> Nil
841
842 -- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at
843 -- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete,
844 -- or update a value in an 'IntMap'. In short : @'lookup' k <$> 'alterF' f k m = f
845 -- ('lookup' k m)@.
846 --
847 -- Example:
848 --
849 -- @
850 -- interactiveAlter :: Int -> IntMap String -> IO (IntMap String)
851 -- interactiveAlter k m = alterF f k m where
852 -- f Nothing -> do
853 -- putStrLn $ show k ++
854 -- " was not found in the map. Would you like to add it?"
855 -- getUserResponse1 :: IO (Maybe String)
856 -- f (Just old) -> do
857 -- putStrLn "The key is currently bound to " ++ show old ++
858 -- ". Would you like to change or delete it?"
859 -- getUserresponse2 :: IO (Maybe String)
860 -- @
861 --
862 -- 'alterF' is the most general operation for working with an individual
863 -- key that may or may not be in a given map.
864 --
865 -- Note: 'alterF' is a flipped version of the 'at' combinator from
866 -- 'Control.Lens.At'.
867 --
868 -- @since 0.5.8
869
870 alterF :: Functor f
871 => (Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a)
872 -- This implementation was stolen from 'Control.Lens.At'.
873 alterF f k m = (<$> f mv) $ \fres ->
874 case fres of
875 Nothing -> maybe m (const (delete k m)) mv
876 Just v' -> insert k v' m
877 where mv = lookup k m
878
879 {--------------------------------------------------------------------
880 Union
881 --------------------------------------------------------------------}
882 -- | The union of a list of maps.
883 --
884 -- > unions [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
885 -- > == fromList [(3, "b"), (5, "a"), (7, "C")]
886 -- > unions [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])]
887 -- > == fromList [(3, "B3"), (5, "A3"), (7, "C")]
888
889 unions :: [IntMap a] -> IntMap a
890 unions xs
891 = foldlStrict union empty xs
892
893 -- | The union of a list of maps, with a combining operation.
894 --
895 -- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
896 -- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")]
897
898 unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a
899 unionsWith f ts
900 = foldlStrict (unionWith f) empty ts
901
902 -- | /O(n+m)/. The (left-biased) union of two maps.
903 -- It prefers the first map when duplicate keys are encountered,
904 -- i.e. (@'union' == 'unionWith' 'const'@).
905 --
906 -- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")]
907
908 union :: IntMap a -> IntMap a -> IntMap a
909 union m1 m2
910 = mergeWithKey' Bin const id id m1 m2
911
912 -- | /O(n+m)/. The union with a combining function.
913 --
914 -- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")]
915
916 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
917 unionWith f m1 m2
918 = unionWithKey (\_ x y -> f x y) m1 m2
919
920 -- | /O(n+m)/. The union with a combining function.
921 --
922 -- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value
923 -- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")]
924
925 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
926 unionWithKey f m1 m2
927 = mergeWithKey' Bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 (f k1 x1 x2)) id id m1 m2
928
929 {--------------------------------------------------------------------
930 Difference
931 --------------------------------------------------------------------}
932 -- | /O(n+m)/. Difference between two maps (based on keys).
933 --
934 -- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b"
935
936 difference :: IntMap a -> IntMap b -> IntMap a
937 difference m1 m2
938 = mergeWithKey (\_ _ _ -> Nothing) id (const Nil) m1 m2
939
940 -- | /O(n+m)/. Difference with a combining function.
941 --
942 -- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing
943 -- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")])
944 -- > == singleton 3 "b:B"
945
946 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
947 differenceWith f m1 m2
948 = differenceWithKey (\_ x y -> f x y) m1 m2
949
950 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
951 -- encountered, the combining function is applied to the key and both values.
952 -- If it returns 'Nothing', the element is discarded (proper set difference).
953 -- If it returns (@'Just' y@), the element is updated with a new value @y@.
954 --
955 -- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing
956 -- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")])
957 -- > == singleton 3 "3:b|B"
958
959 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
960 differenceWithKey f m1 m2
961 = mergeWithKey f id (const Nil) m1 m2
962
963 -- | Remove all the keys in a given set from a map.
964 --
965 -- @
966 -- m `withoutKeys` s = 'filterWithKey' (\k _ -> k `'IntSet.notMember'` s) m
967 -- @
968 --
969 -- @since 0.5.8
970 withoutKeys :: IntMap a -> IntSet.IntSet -> IntMap a
971 withoutKeys = go
972 where
973 go t1@(Bin p1 m1 l1 r1) t2@(IntSet.Bin p2 m2 l2 r2)
974 | shorter m1 m2 = merge1
975 | shorter m2 m1 = merge2
976 | p1 == p2 = bin p1 m1 (go l1 l2) (go r1 r2)
977 | otherwise = t1
978 where
979 merge1 | nomatch p2 p1 m1 = t1
980 | zero p2 m1 = binCheckLeft p1 m1 (go l1 t2) r1
981 | otherwise = binCheckRight p1 m1 l1 (go r1 t2)
982 merge2 | nomatch p1 p2 m2 = t1
983 | zero p1 m2 = bin p2 m2 (go t1 l2) Nil
984 | otherwise = bin p2 m2 Nil (go t1 r2)
985
986 go t1'@(Bin _ _ _ _) t2'@(IntSet.Tip k2' _) = merge t2' k2' t1'
987 where merge t2 k2 t1@(Bin p1 m1 l1 r1) | nomatch k2 p1 m1 = t1
988 | zero k2 m1 = binCheckLeft p1 m1 (merge t2 k2 l1) r1
989 | otherwise = binCheckRight p1 m1 l1 (merge t2 k2 r1)
990 merge _ k2 t1@(Tip k1 _) | k1 == k2 = Nil
991 | otherwise = t1
992 merge _ _ Nil = Nil
993
994 go t1@(Bin _ _ _ _) IntSet.Nil = t1
995
996 go t1'@(Tip k1' _) t2' = merge t1' k1' t2'
997 where merge t1 k1 (IntSet.Bin p2 m2 l2 r2) | nomatch k1 p2 m2 = t1
998 | zero k1 m2 = bin p2 m2 (merge t1 k1 l2) Nil
999 | otherwise = bin p2 m2 Nil (merge t1 k1 r2)
1000 merge t1 k1 (IntSet.Tip k2 _) | k1 == k2 = Nil
1001 | otherwise = t1
1002 merge t1 _ IntSet.Nil = t1
1003
1004 go Nil _ = Nil
1005
1006
1007 {--------------------------------------------------------------------
1008 Intersection
1009 --------------------------------------------------------------------}
1010 -- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys).
1011 --
1012 -- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a"
1013
1014 intersection :: IntMap a -> IntMap b -> IntMap a
1015 intersection m1 m2
1016 = mergeWithKey' bin const (const Nil) (const Nil) m1 m2
1017
1018 -- | /O(n+m)/. The restriction of a map to the keys in a set.
1019 --
1020 -- @
1021 -- m `restrictKeys` s = 'filterWithKey' (\k _ -> k `'IntSet.member'` s) m
1022 -- @
1023 --
1024 -- @since 0.5.8
1025 restrictKeys :: IntMap a -> IntSet.IntSet -> IntMap a
1026 restrictKeys = go
1027 where
1028 go t1@(Bin p1 m1 l1 r1) t2@(IntSet.Bin p2 m2 l2 r2)
1029 | shorter m1 m2 = merge1
1030 | shorter m2 m1 = merge2
1031 | p1 == p2 = bin p1 m1 (go l1 l2) (go r1 r2)
1032 | otherwise = Nil
1033 where
1034 merge1 | nomatch p2 p1 m1 = Nil
1035 | zero p2 m1 = bin p1 m1 (go l1 t2) Nil
1036 | otherwise = bin p1 m1 Nil (go r1 t2)
1037 merge2 | nomatch p1 p2 m2 = Nil
1038 | zero p1 m2 = bin p2 m2 (go t1 l2) Nil
1039 | otherwise = bin p2 m2 Nil (go t1 r2)
1040
1041 go t1'@(Bin _ _ _ _) t2'@(IntSet.Tip k2' _) = merge t2' k2' t1'
1042 where merge t2 k2 (Bin p1 m1 l1 r1) | nomatch k2 p1 m1 = Nil
1043 | zero k2 m1 = bin p1 m1 (merge t2 k2 l1) Nil
1044 | otherwise = bin p1 m1 Nil (merge t2 k2 r1)
1045 merge _ k2 t1@(Tip k1 _) | k1 == k2 = t1
1046 | otherwise = Nil
1047 merge _ _ Nil = Nil
1048
1049 go (Bin _ _ _ _) IntSet.Nil = Nil
1050
1051 go t1'@(Tip k1' _) t2' = merge t1' k1' t2'
1052 where merge t1 k1 (IntSet.Bin p2 m2 l2 r2)
1053 | nomatch k1 p2 m2 = Nil
1054 | zero k1 m2 = bin p2 m2 (merge t1 k1 l2) Nil
1055 | otherwise = bin p2 m2 Nil (merge t1 k1 r2)
1056 merge t1 k1 (IntSet.Tip k2 _) | k1 == k2 = t1
1057 | otherwise = Nil
1058 merge _ _ IntSet.Nil = Nil
1059
1060 go Nil _ = Nil
1061
1062 -- | /O(n+m)/. The intersection with a combining function.
1063 --
1064 -- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA"
1065
1066 intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
1067 intersectionWith f m1 m2
1068 = intersectionWithKey (\_ x y -> f x y) m1 m2
1069
1070 -- | /O(n+m)/. The intersection with a combining function.
1071 --
1072 -- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar
1073 -- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A"
1074
1075 intersectionWithKey :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
1076 intersectionWithKey f m1 m2
1077 = mergeWithKey' bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 (f k1 x1 x2)) (const Nil) (const Nil) m1 m2
1078
1079 {--------------------------------------------------------------------
1080 MergeWithKey
1081 --------------------------------------------------------------------}
1082
1083 -- | /O(n+m)/. A high-performance universal combining function. Using
1084 -- 'mergeWithKey', all combining functions can be defined without any loss of
1085 -- efficiency (with exception of 'union', 'difference' and 'intersection',
1086 -- where sharing of some nodes is lost with 'mergeWithKey').
1087 --
1088 -- Please make sure you know what is going on when using 'mergeWithKey',
1089 -- otherwise you can be surprised by unexpected code growth or even
1090 -- corruption of the data structure.
1091 --
1092 -- When 'mergeWithKey' is given three arguments, it is inlined to the call
1093 -- site. You should therefore use 'mergeWithKey' only to define your custom
1094 -- combining functions. For example, you could define 'unionWithKey',
1095 -- 'differenceWithKey' and 'intersectionWithKey' as
1096 --
1097 -- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2
1098 -- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2
1099 -- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2
1100 --
1101 -- When calling @'mergeWithKey' combine only1 only2@, a function combining two
1102 -- 'IntMap's is created, such that
1103 --
1104 -- * if a key is present in both maps, it is passed with both corresponding
1105 -- values to the @combine@ function. Depending on the result, the key is either
1106 -- present in the result with specified value, or is left out;
1107 --
1108 -- * a nonempty subtree present only in the first map is passed to @only1@ and
1109 -- the output is added to the result;
1110 --
1111 -- * a nonempty subtree present only in the second map is passed to @only2@ and
1112 -- the output is added to the result.
1113 --
1114 -- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/.
1115 -- The values can be modified arbitrarily. Most common variants of @only1@ and
1116 -- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@ or
1117 -- @'filterWithKey' f@ could be used for any @f@.
1118
1119 mergeWithKey :: (Key -> a -> b -> Maybe c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c)
1120 -> IntMap a -> IntMap b -> IntMap c
1121 mergeWithKey f g1 g2 = mergeWithKey' bin combine g1 g2
1122 where -- We use the lambda form to avoid non-exhaustive pattern matches warning.
1123 combine = \(Tip k1 x1) (Tip _k2 x2) -> case f k1 x1 x2 of Nothing -> Nil
1124 Just x -> Tip k1 x
1125 {-# INLINE combine #-}
1126 {-# INLINE mergeWithKey #-}
1127
1128 -- Slightly more general version of mergeWithKey. It differs in the following:
1129 --
1130 -- * the combining function operates on maps instead of keys and values. The
1131 -- reason is to enable sharing in union, difference and intersection.
1132 --
1133 -- * mergeWithKey' is given an equivalent of bin. The reason is that in union*,
1134 -- Bin constructor can be used, because we know both subtrees are nonempty.
1135
1136 mergeWithKey' :: (Prefix -> Mask -> IntMap c -> IntMap c -> IntMap c)
1137 -> (IntMap a -> IntMap b -> IntMap c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c)
1138 -> IntMap a -> IntMap b -> IntMap c
1139 mergeWithKey' bin' f g1 g2 = go
1140 where
1141 go t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
1142 | shorter m1 m2 = merge1
1143 | shorter m2 m1 = merge2
1144 | p1 == p2 = bin' p1 m1 (go l1 l2) (go r1 r2)
1145 | otherwise = maybe_link p1 (g1 t1) p2 (g2 t2)
1146 where
1147 merge1 | nomatch p2 p1 m1 = maybe_link p1 (g1 t1) p2 (g2 t2)
1148 | zero p2 m1 = bin' p1 m1 (go l1 t2) (g1 r1)
1149 | otherwise = bin' p1 m1 (g1 l1) (go r1 t2)
1150 merge2 | nomatch p1 p2 m2 = maybe_link p1 (g1 t1) p2 (g2 t2)
1151 | zero p1 m2 = bin' p2 m2 (go t1 l2) (g2 r2)
1152 | otherwise = bin' p2 m2 (g2 l2) (go t1 r2)
1153
1154 go t1'@(Bin _ _ _ _) t2'@(Tip k2' _) = merge t2' k2' t1'
1155 where merge t2 k2 t1@(Bin p1 m1 l1 r1) | nomatch k2 p1 m1 = maybe_link p1 (g1 t1) k2 (g2 t2)
1156 | zero k2 m1 = bin' p1 m1 (merge t2 k2 l1) (g1 r1)
1157 | otherwise = bin' p1 m1 (g1 l1) (merge t2 k2 r1)
1158 merge t2 k2 t1@(Tip k1 _) | k1 == k2 = f t1 t2
1159 | otherwise = maybe_link k1 (g1 t1) k2 (g2 t2)
1160 merge t2 _ Nil = g2 t2
1161
1162 go t1@(Bin _ _ _ _) Nil = g1 t1
1163
1164 go t1'@(Tip k1' _) t2' = merge t1' k1' t2'
1165 where merge t1 k1 t2@(Bin p2 m2 l2 r2) | nomatch k1 p2 m2 = maybe_link k1 (g1 t1) p2 (g2 t2)
1166 | zero k1 m2 = bin' p2 m2 (merge t1 k1 l2) (g2 r2)
1167 | otherwise = bin' p2 m2 (g2 l2) (merge t1 k1 r2)
1168 merge t1 k1 t2@(Tip k2 _) | k1 == k2 = f t1 t2
1169 | otherwise = maybe_link k1 (g1 t1) k2 (g2 t2)
1170 merge t1 _ Nil = g1 t1
1171
1172 go Nil t2 = g2 t2
1173
1174 maybe_link _ Nil _ t2 = t2
1175 maybe_link _ t1 _ Nil = t1
1176 maybe_link p1 t1 p2 t2 = link p1 t1 p2 t2
1177 {-# INLINE maybe_link #-}
1178 {-# INLINE mergeWithKey' #-}
1179
1180 {--------------------------------------------------------------------
1181 Min\/Max
1182 --------------------------------------------------------------------}
1183
1184 -- | /O(min(n,W))/. Update the value at the minimal key.
1185 --
1186 -- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"3:b"), (5,"a")]
1187 -- > updateMinWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1188
1189 updateMinWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a
1190 updateMinWithKey f t =
1191 case t of Bin p m l r | m < 0 -> binCheckRight p m l (go f r)
1192 _ -> go f t
1193 where
1194 go f' (Bin p m l r) = binCheckLeft p m (go f' l) r
1195 go f' (Tip k y) = case f' k y of
1196 Just y' -> Tip k y'
1197 Nothing -> Nil
1198 go _ Nil = error "updateMinWithKey Nil"
1199
1200 -- | /O(min(n,W))/. Update the value at the maximal key.
1201 --
1202 -- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")]
1203 -- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1204
1205 updateMaxWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a
1206 updateMaxWithKey f t =
1207 case t of Bin p m l r | m < 0 -> binCheckLeft p m (go f l) r
1208 _ -> go f t
1209 where
1210 go f' (Bin p m l r) = binCheckRight p m l (go f' r)
1211 go f' (Tip k y) = case f' k y of
1212 Just y' -> Tip k y'
1213 Nothing -> Nil
1214 go _ Nil = error "updateMaxWithKey Nil"
1215
1216 -- | /O(min(n,W))/. Retrieves the maximal (key,value) pair of the map, and
1217 -- the map stripped of that element, or 'Nothing' if passed an empty map.
1218 --
1219 -- > maxViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((5,"a"), singleton 3 "b")
1220 -- > maxViewWithKey empty == Nothing
1221
1222 maxViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
1223 maxViewWithKey t =
1224 case t of Nil -> Nothing
1225 Bin p m l r | m < 0 -> case go l of (result, l') -> Just (result, binCheckLeft p m l' r)
1226 _ -> Just (go t)
1227 where
1228 go (Bin p m l r) = case go r of (result, r') -> (result, binCheckRight p m l r')
1229 go (Tip k y) = ((k, y), Nil)
1230 go Nil = error "maxViewWithKey Nil"
1231
1232 -- | /O(min(n,W))/. Retrieves the minimal (key,value) pair of the map, and
1233 -- the map stripped of that element, or 'Nothing' if passed an empty map.
1234 --
1235 -- > minViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((3,"b"), singleton 5 "a")
1236 -- > minViewWithKey empty == Nothing
1237
1238 minViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
1239 minViewWithKey t =
1240 case t of Nil -> Nothing
1241 Bin p m l r | m < 0 -> case go r of (result, r') -> Just (result, binCheckRight p m l r')
1242 _ -> Just (go t)
1243 where
1244 go (Bin p m l r) = case go l of (result, l') -> (result, binCheckLeft p m l' r)
1245 go (Tip k y) = ((k, y), Nil)
1246 go Nil = error "minViewWithKey Nil"
1247
1248 -- | /O(min(n,W))/. Update the value at the maximal key.
1249 --
1250 -- > updateMax (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "Xa")]
1251 -- > updateMax (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1252
1253 updateMax :: (a -> Maybe a) -> IntMap a -> IntMap a
1254 updateMax f = updateMaxWithKey (const f)
1255
1256 -- | /O(min(n,W))/. Update the value at the minimal key.
1257 --
1258 -- > updateMin (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "Xb"), (5, "a")]
1259 -- > updateMin (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1260
1261 updateMin :: (a -> Maybe a) -> IntMap a -> IntMap a
1262 updateMin f = updateMinWithKey (const f)
1263
1264 -- Similar to the Arrow instance.
1265 first :: (a -> c) -> (a, b) -> (c, b)
1266 first f (x,y) = (f x,y)
1267
1268 -- | /O(min(n,W))/. Retrieves the maximal key of the map, and the map
1269 -- stripped of that element, or 'Nothing' if passed an empty map.
1270 maxView :: IntMap a -> Maybe (a, IntMap a)
1271 maxView t = liftM (first snd) (maxViewWithKey t)
1272
1273 -- | /O(min(n,W))/. Retrieves the minimal key of the map, and the map
1274 -- stripped of that element, or 'Nothing' if passed an empty map.
1275 minView :: IntMap a -> Maybe (a, IntMap a)
1276 minView t = liftM (first snd) (minViewWithKey t)
1277
1278 -- | /O(min(n,W))/. Delete and find the maximal element.
1279 deleteFindMax :: IntMap a -> ((Key, a), IntMap a)
1280 deleteFindMax = fromMaybe (error "deleteFindMax: empty map has no maximal element") . maxViewWithKey
1281
1282 -- | /O(min(n,W))/. Delete and find the minimal element.
1283 deleteFindMin :: IntMap a -> ((Key, a), IntMap a)
1284 deleteFindMin = fromMaybe (error "deleteFindMin: empty map has no minimal element") . minViewWithKey
1285
1286 -- | /O(min(n,W))/. The minimal key of the map.
1287 findMin :: IntMap a -> (Key, a)
1288 findMin Nil = error $ "findMin: empty map has no minimal element"
1289 findMin (Tip k v) = (k,v)
1290 findMin (Bin _ m l r)
1291 | m < 0 = go r
1292 | otherwise = go l
1293 where go (Tip k v) = (k,v)
1294 go (Bin _ _ l' _) = go l'
1295 go Nil = error "findMax Nil"
1296
1297 -- | /O(min(n,W))/. The maximal key of the map.
1298 findMax :: IntMap a -> (Key, a)
1299 findMax Nil = error $ "findMax: empty map has no maximal element"
1300 findMax (Tip k v) = (k,v)
1301 findMax (Bin _ m l r)
1302 | m < 0 = go l
1303 | otherwise = go r
1304 where go (Tip k v) = (k,v)
1305 go (Bin _ _ _ r') = go r'
1306 go Nil = error "findMax Nil"
1307
1308 -- | /O(min(n,W))/. Delete the minimal key. Returns an empty map if the map is empty.
1309 --
1310 -- Note that this is a change of behaviour for consistency with 'Data.Map.Map' &#8211;
1311 -- versions prior to 0.5 threw an error if the 'IntMap' was already empty.
1312 deleteMin :: IntMap a -> IntMap a
1313 deleteMin = maybe Nil snd . minView
1314
1315 -- | /O(min(n,W))/. Delete the maximal key. Returns an empty map if the map is empty.
1316 --
1317 -- Note that this is a change of behaviour for consistency with 'Data.Map.Map' &#8211;
1318 -- versions prior to 0.5 threw an error if the 'IntMap' was already empty.
1319 deleteMax :: IntMap a -> IntMap a
1320 deleteMax = maybe Nil snd . maxView
1321
1322
1323 {--------------------------------------------------------------------
1324 Submap
1325 --------------------------------------------------------------------}
1326 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
1327 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
1328 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
1329 isProperSubmapOf m1 m2
1330 = isProperSubmapOfBy (==) m1 m2
1331
1332 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
1333 The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
1334 @m1@ and @m2@ are not equal,
1335 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
1336 applied to their respective values. For example, the following
1337 expressions are all 'True':
1338
1339 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1340 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1341
1342 But the following are all 'False':
1343
1344 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
1345 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
1346 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1347 -}
1348 isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
1349 isProperSubmapOfBy predicate t1 t2
1350 = case submapCmp predicate t1 t2 of
1351 LT -> True
1352 _ -> False
1353
1354 submapCmp :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Ordering
1355 submapCmp predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1356 | shorter m1 m2 = GT
1357 | shorter m2 m1 = submapCmpLt
1358 | p1 == p2 = submapCmpEq
1359 | otherwise = GT -- disjoint
1360 where
1361 submapCmpLt | nomatch p1 p2 m2 = GT
1362 | zero p1 m2 = submapCmp predicate t1 l2
1363 | otherwise = submapCmp predicate t1 r2
1364 submapCmpEq = case (submapCmp predicate l1 l2, submapCmp predicate r1 r2) of
1365 (GT,_ ) -> GT
1366 (_ ,GT) -> GT
1367 (EQ,EQ) -> EQ
1368 _ -> LT
1369
1370 submapCmp _ (Bin _ _ _ _) _ = GT
1371 submapCmp predicate (Tip kx x) (Tip ky y)
1372 | (kx == ky) && predicate x y = EQ
1373 | otherwise = GT -- disjoint
1374 submapCmp predicate (Tip k x) t
1375 = case lookup k t of
1376 Just y | predicate x y -> LT
1377 _ -> GT -- disjoint
1378 submapCmp _ Nil Nil = EQ
1379 submapCmp _ Nil _ = LT
1380
1381 -- | /O(n+m)/. Is this a submap?
1382 -- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
1383 isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
1384 isSubmapOf m1 m2
1385 = isSubmapOfBy (==) m1 m2
1386
1387 {- | /O(n+m)/.
1388 The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if
1389 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
1390 applied to their respective values. For example, the following
1391 expressions are all 'True':
1392
1393 > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1394 > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1395 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
1396
1397 But the following are all 'False':
1398
1399 > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)])
1400 > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1401 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
1402 -}
1403 isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
1404 isSubmapOfBy predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1405 | shorter m1 m2 = False
1406 | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy predicate t1 l2
1407 else isSubmapOfBy predicate t1 r2)
1408 | otherwise = (p1==p2) && isSubmapOfBy predicate l1 l2 && isSubmapOfBy predicate r1 r2
1409 isSubmapOfBy _ (Bin _ _ _ _) _ = False
1410 isSubmapOfBy predicate (Tip k x) t = case lookup k t of
1411 Just y -> predicate x y
1412 Nothing -> False
1413 isSubmapOfBy _ Nil _ = True
1414
1415 {--------------------------------------------------------------------
1416 Mapping
1417 --------------------------------------------------------------------}
1418 -- | /O(n)/. Map a function over all values in the map.
1419 --
1420 -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")]
1421
1422 map :: (a -> b) -> IntMap a -> IntMap b
1423 map f = go
1424 where
1425 go (Bin p m l r) = Bin p m (go l) (go r)
1426 go (Tip k x) = Tip k (f x)
1427 go Nil = Nil
1428
1429 #ifdef __GLASGOW_HASKELL__
1430 {-# NOINLINE [1] map #-}
1431 {-# RULES
1432 "map/map" forall f g xs . map f (map g xs) = map (f . g) xs
1433 #-}
1434 #endif
1435 #if __GLASGOW_HASKELL__ >= 709
1436 -- Safe coercions were introduced in 7.8, but did not play well with RULES yet.
1437 {-# RULES
1438 "map/coerce" map coerce = coerce
1439 #-}
1440 #endif
1441
1442 -- | /O(n)/. Map a function over all values in the map.
1443 --
1444 -- > let f key x = (show key) ++ ":" ++ x
1445 -- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")]
1446
1447 mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b
1448 mapWithKey f t
1449 = case t of
1450 Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r)
1451 Tip k x -> Tip k (f k x)
1452 Nil -> Nil
1453
1454 #ifdef __GLASGOW_HASKELL__
1455 {-# NOINLINE [1] mapWithKey #-}
1456 {-# RULES
1457 "mapWithKey/mapWithKey" forall f g xs . mapWithKey f (mapWithKey g xs) =
1458 mapWithKey (\k a -> f k (g k a)) xs
1459 "mapWithKey/map" forall f g xs . mapWithKey f (map g xs) =
1460 mapWithKey (\k a -> f k (g a)) xs
1461 "map/mapWithKey" forall f g xs . map f (mapWithKey g xs) =
1462 mapWithKey (\k a -> f (g k a)) xs
1463 #-}
1464 #endif
1465
1466 -- | /O(n)/.
1467 -- @'traverseWithKey' f s == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@
1468 -- That is, behaves exactly like a regular 'traverse' except that the traversing
1469 -- function also has access to the key associated with a value.
1470 --
1471 -- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(1, 'a'), (5, 'e')]) == Just (fromList [(1, 'b'), (5, 'f')])
1472 -- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(2, 'c')]) == Nothing
1473 traverseWithKey :: Applicative t => (Key -> a -> t b) -> IntMap a -> t (IntMap b)
1474 traverseWithKey f = go
1475 where
1476 go Nil = pure Nil
1477 go (Tip k v) = Tip k <$> f k v
1478 go (Bin p m l r) = Bin p m <$> go l <*> go r
1479 {-# INLINE traverseWithKey #-}
1480
1481 -- | /O(n)/. The function @'mapAccum'@ threads an accumulating
1482 -- argument through the map in ascending order of keys.
1483 --
1484 -- > let f a b = (a ++ b, b ++ "X")
1485 -- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")])
1486
1487 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1488 mapAccum f = mapAccumWithKey (\a' _ x -> f a' x)
1489
1490 -- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating
1491 -- argument through the map in ascending order of keys.
1492 --
1493 -- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X")
1494 -- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")])
1495
1496 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1497 mapAccumWithKey f a t
1498 = mapAccumL f a t
1499
1500 -- | /O(n)/. The function @'mapAccumL'@ threads an accumulating
1501 -- argument through the map in ascending order of keys.
1502 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1503 mapAccumL f a t
1504 = case t of
1505 Bin p m l r -> let (a1,l') = mapAccumL f a l
1506 (a2,r') = mapAccumL f a1 r
1507 in (a2,Bin p m l' r')
1508 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
1509 Nil -> (a,Nil)
1510
1511 -- | /O(n)/. The function @'mapAccumR'@ threads an accumulating
1512 -- argument through the map in descending order of keys.
1513 mapAccumRWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1514 mapAccumRWithKey f a t
1515 = case t of
1516 Bin p m l r -> let (a1,r') = mapAccumRWithKey f a r
1517 (a2,l') = mapAccumRWithKey f a1 l
1518 in (a2,Bin p m l' r')
1519 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
1520 Nil -> (a,Nil)
1521
1522 -- | /O(n*min(n,W))/.
1523 -- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@.
1524 --
1525 -- The size of the result may be smaller if @f@ maps two or more distinct
1526 -- keys to the same new key. In this case the value at the greatest of the
1527 -- original keys is retained.
1528 --
1529 -- > mapKeys (+ 1) (fromList [(5,"a"), (3,"b")]) == fromList [(4, "b"), (6, "a")]
1530 -- > mapKeys (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "c"
1531 -- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c"
1532
1533 mapKeys :: (Key->Key) -> IntMap a -> IntMap a
1534 mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) []
1535
1536 -- | /O(n*min(n,W))/.
1537 -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@.
1538 --
1539 -- The size of the result may be smaller if @f@ maps two or more distinct
1540 -- keys to the same new key. In this case the associated values will be
1541 -- combined using @c@.
1542 --
1543 -- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab"
1544 -- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab"
1545
1546 mapKeysWith :: (a -> a -> a) -> (Key->Key) -> IntMap a -> IntMap a
1547 mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) []
1548
1549 -- | /O(n*min(n,W))/.
1550 -- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@
1551 -- is strictly monotonic.
1552 -- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@.
1553 -- /The precondition is not checked./
1554 -- Semi-formally, we have:
1555 --
1556 -- > and [x < y ==> f x < f y | x <- ls, y <- ls]
1557 -- > ==> mapKeysMonotonic f s == mapKeys f s
1558 -- > where ls = keys s
1559 --
1560 -- This means that @f@ maps distinct original keys to distinct resulting keys.
1561 -- This function has slightly better performance than 'mapKeys'.
1562 --
1563 -- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")]
1564
1565 mapKeysMonotonic :: (Key->Key) -> IntMap a -> IntMap a
1566 mapKeysMonotonic f = fromDistinctAscList . foldrWithKey (\k x xs -> (f k, x) : xs) []
1567
1568 {--------------------------------------------------------------------
1569 Filter
1570 --------------------------------------------------------------------}
1571 -- | /O(n)/. Filter all values that satisfy some predicate.
1572 --
1573 -- > filter (> "a") (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1574 -- > filter (> "x") (fromList [(5,"a"), (3,"b")]) == empty
1575 -- > filter (< "a") (fromList [(5,"a"), (3,"b")]) == empty
1576
1577 filter :: (a -> Bool) -> IntMap a -> IntMap a
1578 filter p m
1579 = filterWithKey (\_ x -> p x) m
1580
1581 -- | /O(n)/. Filter all keys\/values that satisfy some predicate.
1582 --
1583 -- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1584
1585 filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a
1586 filterWithKey predicate t
1587 = case t of
1588 Bin p m l r
1589 -> bin p m (filterWithKey predicate l) (filterWithKey predicate r)
1590 Tip k x
1591 | predicate k x -> t
1592 | otherwise -> Nil
1593 Nil -> Nil
1594
1595 -- | /O(n)/. Partition the map according to some predicate. The first
1596 -- map contains all elements that satisfy the predicate, the second all
1597 -- elements that fail the predicate. See also 'split'.
1598 --
1599 -- > partition (> "a") (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
1600 -- > partition (< "x") (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
1601 -- > partition (> "x") (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
1602
1603 partition :: (a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
1604 partition p m
1605 = partitionWithKey (\_ x -> p x) m
1606
1607 -- | /O(n)/. Partition the map according to some predicate. The first
1608 -- map contains all elements that satisfy the predicate, the second all
1609 -- elements that fail the predicate. See also 'split'.
1610 --
1611 -- > partitionWithKey (\ k _ -> k > 3) (fromList [(5,"a"), (3,"b")]) == (singleton 5 "a", singleton 3 "b")
1612 -- > partitionWithKey (\ k _ -> k < 7) (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
1613 -- > partitionWithKey (\ k _ -> k > 7) (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
1614
1615 partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
1616 partitionWithKey predicate0 t0 = toPair $ go predicate0 t0
1617 where
1618 go predicate t
1619 = case t of
1620 Bin p m l r
1621 -> let (l1 :*: l2) = go predicate l
1622 (r1 :*: r2) = go predicate r
1623 in bin p m l1 r1 :*: bin p m l2 r2
1624 Tip k x
1625 | predicate k x -> (t :*: Nil)
1626 | otherwise -> (Nil :*: t)
1627 Nil -> (Nil :*: Nil)
1628
1629 -- | /O(n)/. Map values and collect the 'Just' results.
1630 --
1631 -- > let f x = if x == "a" then Just "new a" else Nothing
1632 -- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a"
1633
1634 mapMaybe :: (a -> Maybe b) -> IntMap a -> IntMap b
1635 mapMaybe f = mapMaybeWithKey (\_ x -> f x)
1636
1637 -- | /O(n)/. Map keys\/values and collect the 'Just' results.
1638 --
1639 -- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing
1640 -- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3"
1641
1642 mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b
1643 mapMaybeWithKey f (Bin p m l r)
1644 = bin p m (mapMaybeWithKey f l) (mapMaybeWithKey f r)
1645 mapMaybeWithKey f (Tip k x) = case f k x of
1646 Just y -> Tip k y
1647 Nothing -> Nil
1648 mapMaybeWithKey _ Nil = Nil
1649
1650 -- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
1651 --
1652 -- > let f a = if a < "c" then Left a else Right a
1653 -- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1654 -- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")])
1655 -- >
1656 -- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1657 -- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1658
1659 mapEither :: (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
1660 mapEither f m
1661 = mapEitherWithKey (\_ x -> f x) m
1662
1663 -- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
1664 --
1665 -- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a)
1666 -- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1667 -- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")])
1668 -- >
1669 -- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1670 -- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")])
1671
1672 mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
1673 mapEitherWithKey f0 t0 = toPair $ go f0 t0
1674 where
1675 go f (Bin p m l r)
1676 = bin p m l1 r1 :*: bin p m l2 r2
1677 where
1678 (l1 :*: l2) = go f l
1679 (r1 :*: r2) = go f r
1680 go f (Tip k x) = case f k x of
1681 Left y -> (Tip k y :*: Nil)
1682 Right z -> (Nil :*: Tip k z)
1683 go _ Nil = (Nil :*: Nil)
1684
1685 -- | /O(min(n,W))/. The expression (@'split' k map@) is a pair @(map1,map2)@
1686 -- where all keys in @map1@ are lower than @k@ and all keys in
1687 -- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
1688 --
1689 -- > split 2 (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3,"b"), (5,"a")])
1690 -- > split 3 (fromList [(5,"a"), (3,"b")]) == (empty, singleton 5 "a")
1691 -- > split 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
1692 -- > split 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", empty)
1693 -- > split 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], empty)
1694
1695 split :: Key -> IntMap a -> (IntMap a, IntMap a)
1696 split k t =
1697 case t of
1698 Bin _ m l r
1699 | m < 0 -> if k >= 0 -- handle negative numbers.
1700 then case go k l of (lt :*: gt) -> let !lt' = union r lt
1701 in (lt', gt)
1702 else case go k r of (lt :*: gt) -> let !gt' = union gt l
1703 in (lt, gt')
1704 _ -> case go k t of
1705 (lt :*: gt) -> (lt, gt)
1706 where
1707 go k' t'@(Bin p m l r) | nomatch k' p m = if k' > p then t' :*: Nil else Nil :*: t'
1708 | zero k' m = case go k' l of (lt :*: gt) -> lt :*: union gt r
1709 | otherwise = case go k' r of (lt :*: gt) -> union l lt :*: gt
1710 go k' t'@(Tip ky _) | k' > ky = (t' :*: Nil)
1711 | k' < ky = (Nil :*: t')
1712 | otherwise = (Nil :*: Nil)
1713 go _ Nil = (Nil :*: Nil)
1714
1715 -- | /O(min(n,W))/. Performs a 'split' but also returns whether the pivot
1716 -- key was found in the original map.
1717 --
1718 -- > splitLookup 2 (fromList [(5,"a"), (3,"b")]) == (empty, Nothing, fromList [(3,"b"), (5,"a")])
1719 -- > splitLookup 3 (fromList [(5,"a"), (3,"b")]) == (empty, Just "b", singleton 5 "a")
1720 -- > splitLookup 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Nothing, singleton 5 "a")
1721 -- > splitLookup 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Just "a", empty)
1722 -- > splitLookup 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], Nothing, empty)
1723
1724 splitLookup :: Key -> IntMap a -> (IntMap a, Maybe a, IntMap a)
1725 splitLookup k t =
1726 case t of
1727 Bin _ m l r
1728 | m < 0 -> if k >= 0 -- handle negative numbers.
1729 then case go k l of
1730 (lt, fnd, gt) -> let !lt' = union r lt
1731 in (lt', fnd, gt)
1732 else case go k r of
1733 (lt, fnd, gt) -> let !gt' = union gt l
1734 in (lt, fnd, gt')
1735 _ -> go k t
1736 where
1737 go k' t'@(Bin p m l r)
1738 | nomatch k' p m = if k' > p then (t', Nothing, Nil) else (Nil, Nothing, t')
1739 | zero k' m = case go k' l of
1740 (lt, fnd, gt) -> let !gt' = union gt r in (lt, fnd, gt')
1741 | otherwise = case go k' r of
1742 (lt, fnd, gt) -> let !lt' = union l lt in (lt', fnd, gt)
1743 go k' t'@(Tip ky y) | k' > ky = (t', Nothing, Nil)
1744 | k' < ky = (Nil, Nothing, t')
1745 | otherwise = (Nil, Just y, Nil)
1746 go _ Nil = (Nil, Nothing, Nil)
1747
1748 {--------------------------------------------------------------------
1749 Fold
1750 --------------------------------------------------------------------}
1751 -- | /O(n)/. Fold the values in the map using the given right-associative
1752 -- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'elems'@.
1753 --
1754 -- For example,
1755 --
1756 -- > elems map = foldr (:) [] map
1757 --
1758 -- > let f a len = len + (length a)
1759 -- > foldr f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
1760 foldr :: (a -> b -> b) -> b -> IntMap a -> b
1761 foldr f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1762 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1763 | otherwise -> go (go z r) l
1764 _ -> go z t
1765 where
1766 go z' Nil = z'
1767 go z' (Tip _ x) = f x z'
1768 go z' (Bin _ _ l r) = go (go z' r) l
1769 {-# INLINE foldr #-}
1770
1771 -- | /O(n)/. A strict version of 'foldr'. Each application of the operator is
1772 -- evaluated before using the result in the next application. This
1773 -- function is strict in the starting value.
1774 foldr' :: (a -> b -> b) -> b -> IntMap a -> b
1775 foldr' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1776 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1777 | otherwise -> go (go z r) l
1778 _ -> go z t
1779 where
1780 go !z' Nil = z'
1781 go z' (Tip _ x) = f x z'
1782 go z' (Bin _ _ l r) = go (go z' r) l
1783 {-# INLINE foldr' #-}
1784
1785 -- | /O(n)/. Fold the values in the map using the given left-associative
1786 -- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@.
1787 --
1788 -- For example,
1789 --
1790 -- > elems = reverse . foldl (flip (:)) []
1791 --
1792 -- > let f len a = len + (length a)
1793 -- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
1794 foldl :: (a -> b -> a) -> a -> IntMap b -> a
1795 foldl f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1796 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1797 | otherwise -> go (go z l) r
1798 _ -> go z t
1799 where
1800 go z' Nil = z'
1801 go z' (Tip _ x) = f z' x
1802 go z' (Bin _ _ l r) = go (go z' l) r
1803 {-# INLINE foldl #-}
1804
1805 -- | /O(n)/. A strict version of 'foldl'. Each application of the operator is
1806 -- evaluated before using the result in the next application. This
1807 -- function is strict in the starting value.
1808 foldl' :: (a -> b -> a) -> a -> IntMap b -> a
1809 foldl' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1810 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1811 | otherwise -> go (go z l) r
1812 _ -> go z t
1813 where
1814 go !z' Nil = z'
1815 go z' (Tip _ x) = f z' x
1816 go z' (Bin _ _ l r) = go (go z' l) r
1817 {-# INLINE foldl' #-}
1818
1819 -- | /O(n)/. Fold the keys and values in the map using the given right-associative
1820 -- binary operator, such that
1821 -- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
1822 --
1823 -- For example,
1824 --
1825 -- > keys map = foldrWithKey (\k x ks -> k:ks) [] map
1826 --
1827 -- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
1828 -- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)"
1829 foldrWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
1830 foldrWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1831 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1832 | otherwise -> go (go z r) l
1833 _ -> go z t
1834 where
1835 go z' Nil = z'
1836 go z' (Tip kx x) = f kx x z'
1837 go z' (Bin _ _ l r) = go (go z' r) l
1838 {-# INLINE foldrWithKey #-}
1839
1840 -- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is
1841 -- evaluated before using the result in the next application. This
1842 -- function is strict in the starting value.
1843 foldrWithKey' :: (Key -> a -> b -> b) -> b -> IntMap a -> b
1844 foldrWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1845 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1846 | otherwise -> go (go z r) l
1847 _ -> go z t
1848 where
1849 go !z' Nil = z'
1850 go z' (Tip kx x) = f kx x z'
1851 go z' (Bin _ _ l r) = go (go z' r) l
1852 {-# INLINE foldrWithKey' #-}
1853
1854 -- | /O(n)/. Fold the keys and values in the map using the given left-associative
1855 -- binary operator, such that
1856 -- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@.
1857 --
1858 -- For example,
1859 --
1860 -- > keys = reverse . foldlWithKey (\ks k x -> k:ks) []
1861 --
1862 -- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
1863 -- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)"
1864 foldlWithKey :: (a -> Key -> b -> a) -> a -> IntMap b -> a
1865 foldlWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1866 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1867 | otherwise -> go (go z l) r
1868 _ -> go z t
1869 where
1870 go z' Nil = z'
1871 go z' (Tip kx x) = f z' kx x
1872 go z' (Bin _ _ l r) = go (go z' l) r
1873 {-# INLINE foldlWithKey #-}
1874
1875 -- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is
1876 -- evaluated before using the result in the next application. This
1877 -- function is strict in the starting value.
1878 foldlWithKey' :: (a -> Key -> b -> a) -> a -> IntMap b -> a
1879 foldlWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1880 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1881 | otherwise -> go (go z l) r
1882 _ -> go z t
1883 where
1884 go !z' Nil = z'
1885 go z' (Tip kx x) = f z' kx x
1886 go z' (Bin _ _ l r) = go (go z' l) r
1887 {-# INLINE foldlWithKey' #-}
1888
1889 -- | /O(n)/. Fold the keys and values in the map using the given monoid, such that
1890 --
1891 -- @'foldMapWithKey' f = 'Prelude.fold' . 'mapWithKey' f@
1892 --
1893 -- This can be an asymptotically faster than 'foldrWithKey' or 'foldlWithKey' for some monoids.
1894 foldMapWithKey :: Monoid m => (Key -> a -> m) -> IntMap a -> m
1895 foldMapWithKey f = go
1896 where
1897 go Nil = mempty
1898 go (Tip kx x) = f kx x
1899 go (Bin _ _ l r) = go l `mappend` go r
1900 {-# INLINE foldMapWithKey #-}
1901
1902 {--------------------------------------------------------------------
1903 List variations
1904 --------------------------------------------------------------------}
1905 -- | /O(n)/.
1906 -- Return all elements of the map in the ascending order of their keys.
1907 -- Subject to list fusion.
1908 --
1909 -- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"]
1910 -- > elems empty == []
1911
1912 elems :: IntMap a -> [a]
1913 elems = foldr (:) []
1914
1915 -- | /O(n)/. Return all keys of the map in ascending order. Subject to list
1916 -- fusion.
1917 --
1918 -- > keys (fromList [(5,"a"), (3,"b")]) == [3,5]
1919 -- > keys empty == []
1920
1921 keys :: IntMap a -> [Key]
1922 keys = foldrWithKey (\k _ ks -> k : ks) []
1923
1924 -- | /O(n)/. An alias for 'toAscList'. Returns all key\/value pairs in the
1925 -- map in ascending key order. Subject to list fusion.
1926 --
1927 -- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1928 -- > assocs empty == []
1929
1930 assocs :: IntMap a -> [(Key,a)]
1931 assocs = toAscList
1932
1933 -- | /O(n*min(n,W))/. The set of all keys of the map.
1934 --
1935 -- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.IntSet.fromList [3,5]
1936 -- > keysSet empty == Data.IntSet.empty
1937
1938 keysSet :: IntMap a -> IntSet.IntSet
1939 keysSet Nil = IntSet.Nil
1940 keysSet (Tip kx _) = IntSet.singleton kx
1941 keysSet (Bin p m l r)
1942 | m .&. IntSet.suffixBitMask == 0 = IntSet.Bin p m (keysSet l) (keysSet r)
1943 | otherwise = IntSet.Tip (p .&. IntSet.prefixBitMask) (computeBm (computeBm 0 l) r)
1944 where computeBm !acc (Bin _ _ l' r') = computeBm (computeBm acc l') r'
1945 computeBm acc (Tip kx _) = acc .|. IntSet.bitmapOf kx
1946 computeBm _ Nil = error "Data.IntSet.keysSet: Nil"
1947
1948 -- | /O(n)/. Build a map from a set of keys and a function which for each key
1949 -- computes its value.
1950 --
1951 -- > fromSet (\k -> replicate k 'a') (Data.IntSet.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")]
1952 -- > fromSet undefined Data.IntSet.empty == empty
1953
1954 fromSet :: (Key -> a) -> IntSet.IntSet -> IntMap a
1955 fromSet _ IntSet.Nil = Nil
1956 fromSet f (IntSet.Bin p m l r) = Bin p m (fromSet f l) (fromSet f r)
1957 fromSet f (IntSet.Tip kx bm) = buildTree f kx bm (IntSet.suffixBitMask + 1)
1958 where -- This is slightly complicated, as we to convert the dense
1959 -- representation of IntSet into tree representation of IntMap.
1960 --
1961 -- We are given a nonzero bit mask 'bmask' of 'bits' bits with prefix 'prefix'.
1962 -- We split bmask into halves corresponding to left and right subtree.
1963 -- If they are both nonempty, we create a Bin node, otherwise exactly
1964 -- one of them is nonempty and we construct the IntMap from that half.
1965 buildTree g !prefix !bmask bits = case bits of
1966 0 -> Tip prefix (g prefix)
1967 _ -> case intFromNat ((natFromInt bits) `shiftRL` 1) of
1968 bits2 | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 ->
1969 buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2
1970 | (bmask `shiftRL` bits2) .&. ((1 `shiftLL` bits2) - 1) == 0 ->
1971 buildTree g prefix bmask bits2
1972 | otherwise ->
1973 Bin prefix bits2 (buildTree g prefix bmask bits2) (buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2)
1974
1975 {--------------------------------------------------------------------
1976 Lists
1977 --------------------------------------------------------------------}
1978 #if __GLASGOW_HASKELL__ >= 708
1979 instance GHCExts.IsList (IntMap a) where
1980 type Item (IntMap a) = (Key,a)
1981 fromList = fromList
1982 toList = toList
1983 #endif
1984
1985 -- | /O(n)/. Convert the map to a list of key\/value pairs. Subject to list
1986 -- fusion.
1987 --
1988 -- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1989 -- > toList empty == []
1990
1991 toList :: IntMap a -> [(Key,a)]
1992 toList = toAscList
1993
1994 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
1995 -- keys are in ascending order. Subject to list fusion.
1996 --
1997 -- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1998
1999 toAscList :: IntMap a -> [(Key,a)]
2000 toAscList = foldrWithKey (\k x xs -> (k,x):xs) []
2001
2002 -- | /O(n)/. Convert the map to a list of key\/value pairs where the keys
2003 -- are in descending order. Subject to list fusion.
2004 --
2005 -- > toDescList (fromList [(5,"a"), (3,"b")]) == [(5,"a"), (3,"b")]
2006
2007 toDescList :: IntMap a -> [(Key,a)]
2008 toDescList = foldlWithKey (\xs k x -> (k,x):xs) []
2009
2010 -- List fusion for the list generating functions.
2011 #if __GLASGOW_HASKELL__
2012 -- The foldrFB and foldlFB are fold{r,l}WithKey equivalents, used for list fusion.
2013 -- They are important to convert unfused methods back, see mapFB in prelude.
2014 foldrFB :: (Key -> a -> b -> b) -> b -> IntMap a -> b
2015 foldrFB = foldrWithKey
2016 {-# INLINE[0] foldrFB #-}
2017 foldlFB :: (a -> Key -> b -> a) -> a -> IntMap b -> a
2018 foldlFB = foldlWithKey
2019 {-# INLINE[0] foldlFB #-}
2020
2021 -- Inline assocs and toList, so that we need to fuse only toAscList.
2022 {-# INLINE assocs #-}
2023 {-# INLINE toList #-}
2024
2025 -- The fusion is enabled up to phase 2 included. If it does not succeed,
2026 -- convert in phase 1 the expanded elems,keys,to{Asc,Desc}List calls back to
2027 -- elems,keys,to{Asc,Desc}List. In phase 0, we inline fold{lr}FB (which were
2028 -- used in a list fusion, otherwise it would go away in phase 1), and let compiler
2029 -- do whatever it wants with elems,keys,to{Asc,Desc}List -- it was forbidden to
2030 -- inline it before phase 0, otherwise the fusion rules would not fire at all.
2031 {-# NOINLINE[0] elems #-}
2032 {-# NOINLINE[0] keys #-}
2033 {-# NOINLINE[0] toAscList #-}
2034 {-# NOINLINE[0] toDescList #-}
2035 {-# RULES "IntMap.elems" [~1] forall m . elems m = build (\c n -> foldrFB (\_ x xs -> c x xs) n m) #-}
2036 {-# RULES "IntMap.elemsBack" [1] foldrFB (\_ x xs -> x : xs) [] = elems #-}
2037 {-# RULES "IntMap.keys" [~1] forall m . keys m = build (\c n -> foldrFB (\k _ xs -> c k xs) n m) #-}
2038 {-# RULES "IntMap.keysBack" [1] foldrFB (\k _ xs -> k : xs) [] = keys #-}
2039 {-# RULES "IntMap.toAscList" [~1] forall m . toAscList m = build (\c n -> foldrFB (\k x xs -> c (k,x) xs) n m) #-}
2040 {-# RULES "IntMap.toAscListBack" [1] foldrFB (\k x xs -> (k, x) : xs) [] = toAscList #-}
2041 {-# RULES "IntMap.toDescList" [~1] forall m . toDescList m = build (\c n -> foldlFB (\xs k x -> c (k,x) xs) n m) #-}
2042 {-# RULES "IntMap.toDescListBack" [1] foldlFB (\xs k x -> (k, x) : xs) [] = toDescList #-}
2043 #endif
2044
2045
2046 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
2047 --
2048 -- > fromList [] == empty
2049 -- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
2050 -- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
2051
2052 fromList :: [(Key,a)] -> IntMap a
2053 fromList xs
2054 = foldlStrict ins empty xs
2055 where
2056 ins t (k,x) = insert k x t
2057
2058 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
2059 --
2060 -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "ab"), (5, "cba")]
2061 -- > fromListWith (++) [] == empty
2062
2063 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
2064 fromListWith f xs
2065 = fromListWithKey (\_ x y -> f x y) xs
2066
2067 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
2068 --
2069 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
2070 -- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "3:a|b"), (5, "5:c|5:b|a")]
2071 -- > fromListWithKey f [] == empty
2072
2073 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
2074 fromListWithKey f xs
2075 = foldlStrict ins empty xs
2076 where
2077 ins t (k,x) = insertWithKey f k x t
2078
2079 -- | /O(n)/. Build a map from a list of key\/value pairs where
2080 -- the keys are in ascending order.
2081 --
2082 -- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
2083 -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")]
2084
2085 fromAscList :: [(Key,a)] -> IntMap a
2086 fromAscList xs
2087 = fromAscListWithKey (\_ x _ -> x) xs
2088
2089 -- | /O(n)/. Build a map from a list of key\/value pairs where
2090 -- the keys are in ascending order, with a combining function on equal keys.
2091 -- /The precondition (input list is ascending) is not checked./
2092 --
2093 -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
2094
2095 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
2096 fromAscListWith f xs
2097 = fromAscListWithKey (\_ x y -> f x y) xs
2098
2099 -- | /O(n)/. Build a map from a list of key\/value pairs where
2100 -- the keys are in ascending order, with a combining function on equal keys.
2101 -- /The precondition (input list is ascending) is not checked./
2102 --
2103 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
2104 -- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "5:b|a")]
2105
2106 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
2107 fromAscListWithKey _ [] = Nil
2108 fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0)
2109 where
2110 -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
2111 combineEq z [] = [z]
2112 combineEq z@(kz,zz) (x@(kx,xx):xs)
2113 | kx==kz = let yy = f kx xx zz in combineEq (kx,yy) xs
2114 | otherwise = z:combineEq x xs
2115
2116 -- | /O(n)/. Build a map from a list of key\/value pairs where
2117 -- the keys are in ascending order and all distinct.
2118 -- /The precondition (input list is strictly ascending) is not checked./
2119 --
2120 -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
2121
2122 #if __GLASGOW_HASKELL__
2123 fromDistinctAscList :: forall a. [(Key,a)] -> IntMap a
2124 #else
2125 fromDistinctAscList :: [(Key,a)] -> IntMap a
2126 #endif
2127 fromDistinctAscList [] = Nil
2128 fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada
2129 where
2130 work (kx,vx) [] stk = finish kx (Tip kx vx) stk
2131 work (kx,vx) (z@(kz,_):zs) stk = reduce z zs (branchMask kx kz) kx (Tip kx vx) stk
2132
2133 #if __GLASGOW_HASKELL__
2134 reduce :: (Key,a) -> [(Key,a)] -> Mask -> Prefix -> IntMap a -> Stack a -> IntMap a
2135 #endif
2136 reduce z zs _ px tx Nada = work z zs (Push px tx Nada)
2137 reduce z zs m px tx stk@(Push py ty stk') =
2138 let mxy = branchMask px py
2139 pxy = mask px mxy
2140 in if shorter m mxy
2141 then reduce z zs m pxy (Bin pxy mxy ty tx) stk'
2142 else work z zs (Push px tx stk)
2143
2144 finish _ t Nada = t
2145 finish px tx (Push py ty stk) = finish p (link py ty px tx) stk
2146 where m = branchMask px py
2147 p = mask px m
2148
2149 data Stack a = Push {-# UNPACK #-} !Prefix !(IntMap a) !(Stack a) | Nada
2150
2151
2152 {--------------------------------------------------------------------
2153 Eq
2154 --------------------------------------------------------------------}
2155 instance Eq a => Eq (IntMap a) where
2156 t1 == t2 = equal t1 t2
2157 t1 /= t2 = nequal t1 t2
2158
2159 equal :: Eq a => IntMap a -> IntMap a -> Bool
2160 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
2161 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
2162 equal (Tip kx x) (Tip ky y)
2163 = (kx == ky) && (x==y)
2164 equal Nil Nil = True
2165 equal _ _ = False
2166
2167 nequal :: Eq a => IntMap a -> IntMap a -> Bool
2168 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
2169 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
2170 nequal (Tip kx x) (Tip ky y)
2171 = (kx /= ky) || (x/=y)
2172 nequal Nil Nil = False
2173 nequal _ _ = True
2174
2175 {--------------------------------------------------------------------
2176 Ord
2177 --------------------------------------------------------------------}
2178
2179 instance Ord a => Ord (IntMap a) where
2180 compare m1 m2 = compare (toList m1) (toList m2)
2181
2182 {--------------------------------------------------------------------
2183 Functor
2184 --------------------------------------------------------------------}
2185
2186 instance Functor IntMap where
2187 fmap = map
2188
2189 #ifdef __GLASGOW_HASKELL__
2190 a <$ Bin p m l r = Bin p m (a <$ l) (a <$ r)
2191 a <$ Tip k _ = Tip k a
2192 _ <$ Nil = Nil
2193 #endif
2194
2195 {--------------------------------------------------------------------
2196 Show
2197 --------------------------------------------------------------------}
2198
2199 instance Show a => Show (IntMap a) where
2200 showsPrec d m = showParen (d > 10) $
2201 showString "fromList " . shows (toList m)
2202
2203 {--------------------------------------------------------------------
2204 Read
2205 --------------------------------------------------------------------}
2206 instance (Read e) => Read (IntMap e) where
2207 #ifdef __GLASGOW_HASKELL__
2208 readPrec = parens $ prec 10 $ do
2209 Ident "fromList" <- lexP
2210 xs <- readPrec
2211 return (fromList xs)
2212
2213 readListPrec = readListPrecDefault
2214 #else
2215 readsPrec p = readParen (p > 10) $ \ r -> do
2216 ("fromList",s) <- lex r
2217 (xs,t) <- reads s
2218 return (fromList xs,t)
2219 #endif
2220
2221 {--------------------------------------------------------------------
2222 Typeable
2223 --------------------------------------------------------------------}
2224
2225 INSTANCE_TYPEABLE1(IntMap)
2226
2227 {--------------------------------------------------------------------
2228 Helpers
2229 --------------------------------------------------------------------}
2230 {--------------------------------------------------------------------
2231 Link
2232 --------------------------------------------------------------------}
2233 link :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
2234 link p1 t1 p2 t2
2235 | zero p1 m = Bin p m t1 t2
2236 | otherwise = Bin p m t2 t1
2237 where
2238 m = branchMask p1 p2
2239 p = mask p1 m
2240 {-# INLINE link #-}
2241
2242 {--------------------------------------------------------------------
2243 @bin@ assures that we never have empty trees within a tree.
2244 --------------------------------------------------------------------}
2245 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
2246 bin _ _ l Nil = l
2247 bin _ _ Nil r = r
2248 bin p m l r = Bin p m l r
2249 {-# INLINE bin #-}
2250
2251 -- binCheckLeft only checks that the left subtree is non-empty
2252 binCheckLeft :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
2253 binCheckLeft _ _ Nil r = r
2254 binCheckLeft p m l r = Bin p m l r
2255 {-# INLINE binCheckLeft #-}
2256
2257 -- binCheckRight only checks that the right subtree is non-empty
2258 binCheckRight :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
2259 binCheckRight _ _ l Nil = l
2260 binCheckRight p m l r = Bin p m l r
2261 {-# INLINE binCheckRight #-}
2262
2263 {--------------------------------------------------------------------
2264 Endian independent bit twiddling
2265 --------------------------------------------------------------------}
2266 zero :: Key -> Mask -> Bool
2267 zero i m
2268 = (natFromInt i) .&. (natFromInt m) == 0
2269 {-# INLINE zero #-}
2270
2271 nomatch,match :: Key -> Prefix -> Mask -> Bool
2272 nomatch i p m
2273 = (mask i m) /= p
2274 {-# INLINE nomatch #-}
2275
2276 match i p m
2277 = (mask i m) == p
2278 {-# INLINE match #-}
2279
2280 mask :: Key -> Mask -> Prefix
2281 mask i m
2282 = maskW (natFromInt i) (natFromInt m)
2283 {-# INLINE mask #-}
2284
2285
2286 {--------------------------------------------------------------------
2287 Big endian operations
2288 --------------------------------------------------------------------}
2289 maskW :: Nat -> Nat -> Prefix
2290 maskW i m
2291 = intFromNat (i .&. (complement (m-1) `xor` m))
2292 {-# INLINE maskW #-}
2293
2294 shorter :: Mask -> Mask -> Bool
2295 shorter m1 m2
2296 = (natFromInt m1) > (natFromInt m2)
2297 {-# INLINE shorter #-}
2298
2299 branchMask :: Prefix -> Prefix -> Mask
2300 branchMask p1 p2
2301 = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
2302 {-# INLINE branchMask #-}
2303
2304 {--------------------------------------------------------------------
2305 Utilities
2306 --------------------------------------------------------------------}
2307
2308 -- | /O(1)/. Decompose a map into pieces based on the structure of the underlying
2309 -- tree. This function is useful for consuming a map in parallel.
2310 --
2311 -- No guarantee is made as to the sizes of the pieces; an internal, but
2312 -- deterministic process determines this. However, it is guaranteed that the
2313 -- pieces returned will be in ascending order (all elements in the first submap
2314 -- less than all elements in the second, and so on).
2315 --
2316 -- Examples:
2317 --
2318 -- > splitRoot (fromList (zip [1..6::Int] ['a'..])) ==
2319 -- > [fromList [(1,'a'),(2,'b'),(3,'c')],fromList [(4,'d'),(5,'e'),(6,'f')]]
2320 --
2321 -- > splitRoot empty == []
2322 --
2323 -- Note that the current implementation does not return more than two submaps,
2324 -- but you should not depend on this behaviour because it can change in the
2325 -- future without notice.
2326 splitRoot :: IntMap a -> [IntMap a]
2327 splitRoot orig =
2328 case orig of
2329 Nil -> []
2330 x@(Tip _ _) -> [x]
2331 Bin _ m l r | m < 0 -> [r, l]
2332 | otherwise -> [l, r]
2333 {-# INLINE splitRoot #-}
2334
2335
2336 {--------------------------------------------------------------------
2337 Debugging
2338 --------------------------------------------------------------------}
2339 -- | /O(n)/. Show the tree that implements the map. The tree is shown
2340 -- in a compressed, hanging format.
2341 showTree :: Show a => IntMap a -> String
2342 showTree s
2343 = showTreeWith True False s
2344
2345
2346 {- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows
2347 the tree that implements the map. If @hang@ is
2348 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
2349 @wide@ is 'True', an extra wide version is shown.
2350 -}
2351 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
2352 showTreeWith hang wide t
2353 | hang = (showsTreeHang wide [] t) ""
2354 | otherwise = (showsTree wide [] [] t) ""
2355
2356 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
2357 showsTree wide lbars rbars t
2358 = case t of
2359 Bin p m l r
2360 -> showsTree wide (withBar rbars) (withEmpty rbars) r .
2361 showWide wide rbars .
2362 showsBars lbars . showString (showBin p m) . showString "\n" .
2363 showWide wide lbars .
2364 showsTree wide (withEmpty lbars) (withBar lbars) l
2365 Tip k x
2366 -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n"
2367 Nil -> showsBars lbars . showString "|\n"
2368
2369 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
2370 showsTreeHang wide bars t
2371 = case t of
2372 Bin p m l r
2373 -> showsBars bars . showString (showBin p m) . showString "\n" .
2374 showWide wide bars .
2375 showsTreeHang wide (withBar bars) l .
2376 showWide wide bars .
2377 showsTreeHang wide (withEmpty bars) r
2378 Tip k x
2379 -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n"
2380 Nil -> showsBars bars . showString "|\n"
2381
2382 showBin :: Prefix -> Mask -> String
2383 showBin _ _
2384 = "*" -- ++ show (p,m)
2385
2386 showWide :: Bool -> [String] -> String -> String
2387 showWide wide bars
2388 | wide = showString (concat (reverse bars)) . showString "|\n"
2389 | otherwise = id
2390
2391 showsBars :: [String] -> ShowS
2392 showsBars bars
2393 = case bars of
2394 [] -> id
2395 _ -> showString (concat (reverse (tail bars))) . showString node
2396
2397 node :: String
2398 node = "+--"
2399
2400 withBar, withEmpty :: [String] -> [String]
2401 withBar bars = "| ":bars
2402 withEmpty bars = " ":bars