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