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