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