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