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