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