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