Merge pull request #219 from treeowl/intmap-binbetter
[packages/containers.git] / Data / IntMap / Base.hs
1 {-# LANGUAGE CPP #-}
2 {-# LANGUAGE BangPatterns #-}
3 #if __GLASGOW_HASKELL__
4 {-# LANGUAGE MagicHash, DeriveDataTypeable, StandaloneDeriving #-}
5 #endif
6 #if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703
7 {-# LANGUAGE Trustworthy #-}
8 #endif
9 {-# LANGUAGE ScopedTypeVariables #-}
10 #if __GLASGOW_HASKELL__ >= 708
11 {-# LANGUAGE TypeFamilies #-}
12 #endif
13
14 #include "containers.h"
15
16 -----------------------------------------------------------------------------
17 -- |
18 -- Module : Data.IntMap.Base
19 -- Copyright : (c) Daan Leijen 2002
20 -- (c) Andriy Palamarchuk 2008
21 -- License : BSD-style
22 -- Maintainer : libraries@haskell.org
23 -- Stability : provisional
24 -- Portability : portable
25 --
26 -- This defines the data structures and core (hidden) manipulations
27 -- on representations.
28 -----------------------------------------------------------------------------
29
30 -- [Note: INLINE bit fiddling]
31 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~
32 -- It is essential that the bit fiddling functions like mask, zero, branchMask
33 -- etc are inlined. If they do not, the memory allocation skyrockets. The GHC
34 -- usually gets it right, but it is disastrous if it does not. Therefore we
35 -- explicitly mark these functions INLINE.
36
37
38 -- [Note: Local 'go' functions and capturing]
39 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
40 -- Care must be taken when using 'go' function which captures an argument.
41 -- Sometimes (for example when the argument is passed to a data constructor,
42 -- as in insert), GHC heap-allocates more than necessary. Therefore C-- code
43 -- must be checked for increased allocation when creating and modifying such
44 -- functions.
45
46
47 -- [Note: Order of constructors]
48 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
49 -- The order of constructors of IntMap matters when considering performance.
50 -- Currently in GHC 7.0, when type has 3 constructors, they are matched from
51 -- the first to the last -- the best performance is achieved when the
52 -- constructors are ordered by frequency.
53 -- On GHC 7.0, reordering constructors from Nil | Tip | Bin to Bin | Tip | Nil
54 -- improves the benchmark by circa 10%.
55
56 module Data.IntMap.Base (
57 -- * Map type
58 IntMap(..), Key -- instance Eq,Show
59
60 -- * Operators
61 , (!), (\\)
62
63 -- * Query
64 , null
65 , size
66 , member
67 , notMember
68 , lookup
69 , findWithDefault
70 , lookupLT
71 , lookupGT
72 , lookupLE
73 , lookupGE
74
75 -- * Construction
76 , empty
77 , singleton
78
79 -- ** Insertion
80 , insert
81 , insertWith
82 , insertWithKey
83 , insertLookupWithKey
84
85 -- ** Delete\/Update
86 , delete
87 , adjust
88 , adjustWithKey
89 , update
90 , updateWithKey
91 , updateLookupWithKey
92 , alter
93
94 -- * Combine
95
96 -- ** Union
97 , union
98 , unionWith
99 , unionWithKey
100 , unions
101 , unionsWith
102
103 -- ** Difference
104 , difference
105 , differenceWith
106 , differenceWithKey
107
108 -- ** Intersection
109 , intersection
110 , intersectionWith
111 , intersectionWithKey
112
113 -- ** Universal combining function
114 , mergeWithKey
115 , mergeWithKey'
116
117 -- * Traversal
118 -- ** Map
119 , map
120 , mapWithKey
121 , traverseWithKey
122 , mapAccum
123 , mapAccumWithKey
124 , mapAccumRWithKey
125 , mapKeys
126 , mapKeysWith
127 , mapKeysMonotonic
128
129 -- * Folds
130 , foldr
131 , foldl
132 , foldrWithKey
133 , foldlWithKey
134 , foldMapWithKey
135
136 -- ** Strict folds
137 , foldr'
138 , foldl'
139 , foldrWithKey'
140 , foldlWithKey'
141
142 -- * Conversion
143 , elems
144 , keys
145 , assocs
146 , keysSet
147 , fromSet
148
149 -- ** Lists
150 , toList
151 , fromList
152 , fromListWith
153 , fromListWithKey
154
155 -- ** Ordered lists
156 , toAscList
157 , toDescList
158 , fromAscList
159 , fromAscListWith
160 , fromAscListWithKey
161 , fromDistinctAscList
162
163 -- * Filter
164 , filter
165 , filterWithKey
166 , partition
167 , partitionWithKey
168
169 , mapMaybe
170 , mapMaybeWithKey
171 , mapEither
172 , mapEitherWithKey
173
174 , split
175 , splitLookup
176 , splitRoot
177
178 -- * Submap
179 , isSubmapOf, isSubmapOfBy
180 , isProperSubmapOf, isProperSubmapOfBy
181
182 -- * Min\/Max
183 , findMin
184 , findMax
185 , deleteMin
186 , deleteMax
187 , deleteFindMin
188 , deleteFindMax
189 , updateMin
190 , updateMax
191 , updateMinWithKey
192 , updateMaxWithKey
193 , minView
194 , maxView
195 , minViewWithKey
196 , maxViewWithKey
197
198 -- * Debugging
199 , showTree
200 , showTreeWith
201
202 -- * Internal types
203 , Mask, Prefix, Nat
204
205 -- * Utility
206 , natFromInt
207 , intFromNat
208 , link
209 , bin
210 , binCheckLeft
211 , binCheckRight
212 , zero
213 , nomatch
214 , match
215 , mask
216 , maskW
217 , shorter
218 , branchMask
219 , highestBitMask
220 ) where
221
222 #if !(MIN_VERSION_base(4,8,0))
223 import Control.Applicative (Applicative(pure, (<*>)), (<$>))
224 import Data.Monoid (Monoid(..))
225 import Data.Traversable (Traversable(traverse))
226 import Data.Word (Word)
227 #endif
228 #if MIN_VERSION_base(4,9,0)
229 import Data.Semigroup (Semigroup((<>), stimes), stimesIdempotentMonoid)
230 #endif
231
232 import Control.DeepSeq (NFData(rnf))
233 import Control.Monad (liftM)
234 import Data.Bits
235 import qualified Data.Foldable as Foldable
236 import Data.Maybe (fromMaybe)
237 import Data.Typeable
238 import Prelude hiding (lookup, map, filter, foldr, foldl, null)
239
240 import Data.IntSet.Base (Key)
241 import qualified Data.IntSet.Base as IntSet
242 import Data.Utils.BitUtil
243 import Data.Utils.StrictFold
244 import Data.Utils.StrictPair
245
246 #if __GLASGOW_HASKELL__
247 import Data.Data (Data(..), Constr, mkConstr, constrIndex, Fixity(Prefix),
248 DataType, mkDataType)
249 import GHC.Exts (build)
250 #if __GLASGOW_HASKELL__ >= 708
251 import qualified GHC.Exts as GHCExts
252 #endif
253 import Text.Read
254 #endif
255 #if __GLASGOW_HASKELL__ >= 709
256 import Data.Coerce
257 #endif
258
259
260 -- A "Nat" is a natural machine word (an unsigned Int)
261 type Nat = Word
262
263 natFromInt :: Key -> Nat
264 natFromInt = fromIntegral
265 {-# INLINE natFromInt #-}
266
267 intFromNat :: Nat -> Key
268 intFromNat = fromIntegral
269 {-# INLINE intFromNat #-}
270
271 {--------------------------------------------------------------------
272 Types
273 --------------------------------------------------------------------}
274
275
276 -- | A map of integers to values @a@.
277
278 -- See Note: Order of constructors
279 data IntMap a = Bin {-# UNPACK #-} !Prefix
280 {-# UNPACK #-} !Mask
281 !(IntMap a)
282 !(IntMap a)
283 | Tip {-# UNPACK #-} !Key a
284 | Nil
285
286 type Prefix = Int
287 type Mask = Int
288
289 {--------------------------------------------------------------------
290 Operators
291 --------------------------------------------------------------------}
292
293 -- | /O(min(n,W))/. Find the value at a key.
294 -- Calls 'error' when the element can not be found.
295 --
296 -- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map
297 -- > fromList [(5,'a'), (3,'b')] ! 5 == 'a'
298
299 (!) :: IntMap a -> Key -> a
300 (!) m k = find k m
301
302 -- | Same as 'difference'.
303 (\\) :: IntMap a -> IntMap b -> IntMap a
304 m1 \\ m2 = difference m1 m2
305
306 infixl 9 \\{-This comment teaches CPP correct behaviour -}
307
308 {--------------------------------------------------------------------
309 Types
310 --------------------------------------------------------------------}
311
312 instance Monoid (IntMap a) where
313 mempty = empty
314 mconcat = unions
315 #if !(MIN_VERSION_base(4,9,0))
316 mappend = union
317 #else
318 mappend = (<>)
319
320 instance Semigroup (IntMap a) where
321 (<>) = union
322 stimes = stimesIdempotentMonoid
323 #endif
324
325 instance Foldable.Foldable IntMap where
326 fold = go
327 where go Nil = mempty
328 go (Tip _ v) = v
329 go (Bin _ _ l r) = go l `mappend` go r
330 {-# INLINABLE fold #-}
331 foldr = foldr
332 {-# INLINE foldr #-}
333 foldl = foldl
334 {-# INLINE foldl #-}
335 foldMap f t = go t
336 where go Nil = mempty
337 go (Tip _ v) = f v
338 go (Bin _ _ l r) = go l `mappend` go r
339 {-# INLINE foldMap #-}
340
341 #if MIN_VERSION_base(4,6,0)
342 foldl' = foldl'
343 {-# INLINE foldl' #-}
344 foldr' = foldr'
345 {-# INLINE foldr' #-}
346 #endif
347 #if MIN_VERSION_base(4,8,0)
348 length = size
349 {-# INLINE length #-}
350 null = null
351 {-# INLINE null #-}
352 toList = elems -- NB: Foldable.toList /= IntMap.toList
353 {-# INLINE toList #-}
354 elem = go
355 where go !_ Nil = False
356 go x (Tip _ y) = x == y
357 go x (Bin _ _ l r) = go x l || go x r
358 {-# INLINABLE elem #-}
359 maximum = start
360 where start Nil = error "IntMap.Foldable.maximum: called with empty map"
361 start (Tip _ y) = y
362 start (Bin _ _ l r) = go (start l) r
363
364 go !m Nil = m
365 go m (Tip _ y) = max m y
366 go m (Bin _ _ l r) = go (go m l) r
367 {-# INLINABLE maximum #-}
368 minimum = start
369 where start Nil = error "IntMap.Foldable.minimum: called with empty map"
370 start (Tip _ y) = y
371 start (Bin _ _ l r) = go (start l) r
372
373 go !m Nil = m
374 go m (Tip _ y) = min m y
375 go m (Bin _ _ l r) = go (go m l) r
376 {-# INLINABLE minimum #-}
377 sum = foldl' (+) 0
378 {-# INLINABLE sum #-}
379 product = foldl' (*) 1
380 {-# INLINABLE product #-}
381 #endif
382
383 instance Traversable IntMap where
384 traverse f = traverseWithKey (\_ -> f)
385 {-# INLINE traverse #-}
386
387 instance NFData a => NFData (IntMap a) where
388 rnf Nil = ()
389 rnf (Tip _ v) = rnf v
390 rnf (Bin _ _ l r) = rnf l `seq` rnf r
391
392 #if __GLASGOW_HASKELL__
393
394 {--------------------------------------------------------------------
395 A Data instance
396 --------------------------------------------------------------------}
397
398 -- This instance preserves data abstraction at the cost of inefficiency.
399 -- We provide limited reflection services for the sake of data abstraction.
400
401 instance Data a => Data (IntMap a) where
402 gfoldl f z im = z fromList `f` (toList im)
403 toConstr _ = fromListConstr
404 gunfold k z c = case constrIndex c of
405 1 -> k (z fromList)
406 _ -> error "gunfold"
407 dataTypeOf _ = intMapDataType
408 dataCast1 f = gcast1 f
409
410 fromListConstr :: Constr
411 fromListConstr = mkConstr intMapDataType "fromList" [] Prefix
412
413 intMapDataType :: DataType
414 intMapDataType = mkDataType "Data.IntMap.Base.IntMap" [fromListConstr]
415
416 #endif
417
418 {--------------------------------------------------------------------
419 Query
420 --------------------------------------------------------------------}
421 -- | /O(1)/. Is the map empty?
422 --
423 -- > Data.IntMap.null (empty) == True
424 -- > Data.IntMap.null (singleton 1 'a') == False
425
426 null :: IntMap a -> Bool
427 null Nil = True
428 null _ = False
429 {-# INLINE null #-}
430
431 -- | /O(n)/. Number of elements in the map.
432 --
433 -- > size empty == 0
434 -- > size (singleton 1 'a') == 1
435 -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3
436 size :: IntMap a -> Int
437 size (Bin _ _ l r) = size l + size r
438 size (Tip _ _) = 1
439 size Nil = 0
440
441 -- | /O(min(n,W))/. Is the key a member of the map?
442 --
443 -- > member 5 (fromList [(5,'a'), (3,'b')]) == True
444 -- > member 1 (fromList [(5,'a'), (3,'b')]) == False
445
446 -- See Note: Local 'go' functions and capturing]
447 member :: Key -> IntMap a -> Bool
448 member !k = go
449 where
450 go (Bin p m l r) | nomatch k p m = False
451 | zero k m = go l
452 | otherwise = go r
453 go (Tip kx _) = k == kx
454 go Nil = False
455
456 -- | /O(min(n,W))/. Is the key not a member of the map?
457 --
458 -- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False
459 -- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True
460
461 notMember :: Key -> IntMap a -> Bool
462 notMember k m = not $ member k m
463
464 -- | /O(min(n,W))/. Lookup the value at a key in the map. See also 'Data.Map.lookup'.
465
466 -- See Note: Local 'go' functions and capturing]
467 lookup :: Key -> IntMap a -> Maybe a
468 lookup !k = go
469 where
470 go (Bin p m l r) | nomatch k p m = Nothing
471 | zero k m = go l
472 | otherwise = go r
473 go (Tip kx x) | k == kx = Just x
474 | otherwise = Nothing
475 go Nil = Nothing
476
477
478 -- See Note: Local 'go' functions and capturing]
479 find :: Key -> IntMap a -> a
480 find !k = go
481 where
482 go (Bin p m l r) | nomatch k p m = not_found
483 | zero k m = go l
484 | otherwise = go r
485 go (Tip kx x) | k == kx = x
486 | otherwise = not_found
487 go Nil = not_found
488
489 not_found = error ("IntMap.!: key " ++ show k ++ " is not an element of the map")
490
491 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
492 -- returns the value at key @k@ or returns @def@ when the key is not an
493 -- element of the map.
494 --
495 -- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x'
496 -- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a'
497
498 -- See Note: Local 'go' functions and capturing]
499 findWithDefault :: a -> Key -> IntMap a -> a
500 findWithDefault def !k = go
501 where
502 go (Bin p m l r) | nomatch k p m = def
503 | zero k m = go l
504 | otherwise = go r
505 go (Tip kx x) | k == kx = x
506 | otherwise = def
507 go Nil = def
508
509 -- | /O(log n)/. Find largest key smaller than the given one and return the
510 -- corresponding (key, value) pair.
511 --
512 -- > lookupLT 3 (fromList [(3,'a'), (5,'b')]) == Nothing
513 -- > lookupLT 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
514
515 -- See Note: Local 'go' functions and capturing.
516 lookupLT :: Key -> IntMap a -> Maybe (Key, a)
517 lookupLT !k t = case t of
518 Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r
519 _ -> go Nil t
520 where
521 go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMax def else unsafeFindMax r
522 | zero k m = go def l
523 | otherwise = go l r
524 go def (Tip ky y) | k <= ky = unsafeFindMax def
525 | otherwise = Just (ky, y)
526 go def Nil = unsafeFindMax def
527
528 -- | /O(log n)/. Find smallest key greater than the given one and return the
529 -- corresponding (key, value) pair.
530 --
531 -- > lookupGT 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
532 -- > lookupGT 5 (fromList [(3,'a'), (5,'b')]) == Nothing
533
534 -- See Note: Local 'go' functions and capturing.
535 lookupGT :: Key -> IntMap a -> Maybe (Key, a)
536 lookupGT !k t = case t of
537 Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r
538 _ -> go Nil t
539 where
540 go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMin l else unsafeFindMin def
541 | zero k m = go r l
542 | otherwise = go def r
543 go def (Tip ky y) | k >= ky = unsafeFindMin def
544 | otherwise = Just (ky, y)
545 go def Nil = unsafeFindMin def
546
547 -- | /O(log n)/. Find largest key smaller or equal to the given one and return
548 -- the corresponding (key, value) pair.
549 --
550 -- > lookupLE 2 (fromList [(3,'a'), (5,'b')]) == Nothing
551 -- > lookupLE 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
552 -- > lookupLE 5 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
553
554 -- See Note: Local 'go' functions and capturing.
555 lookupLE :: Key -> IntMap a -> Maybe (Key, a)
556 lookupLE !k t = case t of
557 Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r
558 _ -> go Nil t
559 where
560 go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMax def else unsafeFindMax r
561 | zero k m = go def l
562 | otherwise = go l r
563 go def (Tip ky y) | k < ky = unsafeFindMax def
564 | otherwise = Just (ky, y)
565 go def Nil = unsafeFindMax def
566
567 -- | /O(log n)/. Find smallest key greater or equal to the given one and return
568 -- the corresponding (key, value) pair.
569 --
570 -- > lookupGE 3 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
571 -- > lookupGE 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
572 -- > lookupGE 6 (fromList [(3,'a'), (5,'b')]) == Nothing
573
574 -- See Note: Local 'go' functions and capturing.
575 lookupGE :: Key -> IntMap a -> Maybe (Key, a)
576 lookupGE !k t = case t of
577 Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r
578 _ -> go Nil t
579 where
580 go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMin l else unsafeFindMin def
581 | zero k m = go r l
582 | otherwise = go def r
583 go def (Tip ky y) | k > ky = unsafeFindMin def
584 | otherwise = Just (ky, y)
585 go def Nil = unsafeFindMin def
586
587
588 -- Helper function for lookupGE and lookupGT. It assumes that if a Bin node is
589 -- given, it has m > 0.
590 unsafeFindMin :: IntMap a -> Maybe (Key, a)
591 unsafeFindMin Nil = Nothing
592 unsafeFindMin (Tip ky y) = Just (ky, y)
593 unsafeFindMin (Bin _ _ l _) = unsafeFindMin l
594
595 -- Helper function for lookupLE and lookupLT. It assumes that if a Bin node is
596 -- given, it has m > 0.
597 unsafeFindMax :: IntMap a -> Maybe (Key, a)
598 unsafeFindMax Nil = Nothing
599 unsafeFindMax (Tip ky y) = Just (ky, y)
600 unsafeFindMax (Bin _ _ _ r) = unsafeFindMax r
601
602 {--------------------------------------------------------------------
603 Construction
604 --------------------------------------------------------------------}
605 -- | /O(1)/. The empty map.
606 --
607 -- > empty == fromList []
608 -- > size empty == 0
609
610 empty :: IntMap a
611 empty
612 = Nil
613 {-# INLINE empty #-}
614
615 -- | /O(1)/. A map of one element.
616 --
617 -- > singleton 1 'a' == fromList [(1, 'a')]
618 -- > size (singleton 1 'a') == 1
619
620 singleton :: Key -> a -> IntMap a
621 singleton k x
622 = Tip k x
623 {-# INLINE singleton #-}
624
625 {--------------------------------------------------------------------
626 Insert
627 --------------------------------------------------------------------}
628 -- | /O(min(n,W))/. Insert a new key\/value pair in the map.
629 -- If the key is already present in the map, the associated value is
630 -- replaced with the supplied value, i.e. 'insert' is equivalent to
631 -- @'insertWith' 'const'@.
632 --
633 -- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')]
634 -- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')]
635 -- > insert 5 'x' empty == singleton 5 'x'
636
637 insert :: Key -> a -> IntMap a -> IntMap a
638 insert !k x t@(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 insert k x t@(Tip ky _)
643 | k==ky = Tip k x
644 | otherwise = link k (Tip k x) ky t
645 insert k x 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@(Bin p m l r)
675 | nomatch k p m = link k (Tip k x) p t
676 | zero k m = Bin p m (insertWithKey f k x l) r
677 | otherwise = Bin p m l (insertWithKey f k x r)
678 insertWithKey f k x t@(Tip ky y)
679 | k == ky = Tip k (f k x y)
680 | otherwise = link k (Tip k x) ky t
681 insertWithKey _ k x Nil = Tip k x
682
683 -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
684 -- is a pair where the first element is equal to (@'lookup' k map@)
685 -- and the second element equal to (@'insertWithKey' f k x map@).
686 --
687 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
688 -- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")])
689 -- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")])
690 -- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx")
691 --
692 -- This is how to define @insertLookup@ using @insertLookupWithKey@:
693 --
694 -- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t
695 -- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")])
696 -- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")])
697
698 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
699 insertLookupWithKey f !k x t@(Bin p m l r)
700 | nomatch k p m = (Nothing,link k (Tip k x) p t)
701 | zero k m = let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r)
702 | otherwise = let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r')
703 insertLookupWithKey f k x t@(Tip ky y)
704 | k == ky = (Just y,Tip k (f k x y))
705 | otherwise = (Nothing,link k (Tip k x) ky t)
706 insertLookupWithKey _ k x Nil = (Nothing,Tip k x)
707
708
709 {--------------------------------------------------------------------
710 Deletion
711 --------------------------------------------------------------------}
712 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
713 -- a member of the map, the original map is returned.
714 --
715 -- > delete 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
716 -- > delete 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
717 -- > delete 5 empty == empty
718
719 delete :: Key -> IntMap a -> IntMap a
720 delete !k t@(Bin p m l r)
721 | nomatch k p m = t
722 | zero k m = binCheckLeft p m (delete k l) r
723 | otherwise = binCheckRight p m l (delete k r)
724 delete k t@(Tip ky _)
725 | k == ky = Nil
726 | otherwise = t
727 delete _k Nil = Nil
728
729 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
730 -- a member of the map, the original map is returned.
731 --
732 -- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
733 -- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
734 -- > adjust ("new " ++) 7 empty == empty
735
736 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
737 adjust f k m
738 = adjustWithKey (\_ x -> f x) k m
739
740 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
741 -- a member of the map, the original map is returned.
742 --
743 -- > let f key x = (show key) ++ ":new " ++ x
744 -- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
745 -- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
746 -- > adjustWithKey f 7 empty == empty
747
748 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
749 adjustWithKey f !k t@(Bin p m l r)
750 | nomatch k p m = t
751 | zero k m = Bin p m (adjustWithKey f k l) r
752 | otherwise = Bin p m l (adjustWithKey f k r)
753 adjustWithKey f k t@(Tip ky y)
754 | k == ky = Tip ky (f k y)
755 | otherwise = t
756 adjustWithKey _ _ Nil = Nil
757
758
759 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
760 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
761 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
762 --
763 -- > let f x = if x == "a" then Just "new a" else Nothing
764 -- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
765 -- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
766 -- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
767
768 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
769 update f
770 = updateWithKey (\_ x -> f x)
771
772 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
773 -- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
774 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
775 --
776 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
777 -- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
778 -- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
779 -- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
780
781 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
782 updateWithKey f !k t@(Bin p m l r)
783 | nomatch k p m = t
784 | zero k m = binCheckLeft p m (updateWithKey f k l) r
785 | otherwise = binCheckRight p m l (updateWithKey f k r)
786 updateWithKey f k t@(Tip ky y)
787 | k == ky = case (f k y) of
788 Just y' -> Tip ky y'
789 Nothing -> Nil
790 | otherwise = t
791 updateWithKey _ _ Nil = Nil
792
793 -- | /O(min(n,W))/. Lookup and update.
794 -- The function returns original value, if it is updated.
795 -- This is different behavior than 'Data.Map.updateLookupWithKey'.
796 -- Returns the original key value if the map entry is deleted.
797 --
798 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
799 -- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:new a")])
800 -- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")])
801 -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a")
802
803 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
804 updateLookupWithKey f !k t@(Bin p m l r)
805 | nomatch k p m = (Nothing,t)
806 | zero k m = let !(found,l') = updateLookupWithKey f k l in (found,binCheckLeft p m l' r)
807 | otherwise = let !(found,r') = updateLookupWithKey f k r in (found,binCheckRight p m l r')
808 updateLookupWithKey f k t@(Tip ky y)
809 | k==ky = case (f k y) of
810 Just y' -> (Just y,Tip ky y')
811 Nothing -> (Just y,Nil)
812 | otherwise = (Nothing,t)
813 updateLookupWithKey _ _ Nil = (Nothing,Nil)
814
815
816
817 -- | /O(min(n,W))/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
818 -- 'alter' can be used to insert, delete, or update a value in an 'IntMap'.
819 -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
820 alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
821 alter f !k t@(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 = binCheckLeft p m (alter f k l) r
826 | otherwise = binCheckRight p m l (alter f k r)
827 alter f k t@(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 alter f k 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 -> binCheckRight p m l (go f r)
1065 _ -> go f t
1066 where
1067 go f' (Bin p m l r) = binCheckLeft 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 -> binCheckLeft p m (go f l) r
1081 _ -> go f t
1082 where
1083 go f' (Bin p m l r) = binCheckRight 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, binCheckLeft p m l' r)
1099 _ -> Just (go t)
1100 where
1101 go (Bin p m l r) = case go r of (result, r') -> (result, binCheckRight 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, binCheckRight p m l r')
1115 _ -> Just (go t)
1116 where
1117 go (Bin p m l r) = case go l of (result, l') -> (result, binCheckLeft 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', gt)
1575 else case go k r of (lt :*: gt) -> let !gt' = union gt l
1576 in (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', fnd, gt)
1605 else case go k r of
1606 (lt, fnd, gt) -> let !gt' = union gt l
1607 in (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 (lt, fnd, gt')
1614 | otherwise = case go k' r of
1615 (lt, fnd, gt) -> let !lt' = union l lt in (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 go !z' Nil = z'
1654 go z' (Tip _ x) = f x z'
1655 go z' (Bin _ _ l r) = go (go z' r) l
1656 {-# INLINE foldr' #-}
1657
1658 -- | /O(n)/. Fold the values in the map using the given left-associative
1659 -- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@.
1660 --
1661 -- For example,
1662 --
1663 -- > elems = reverse . foldl (flip (:)) []
1664 --
1665 -- > let f len a = len + (length a)
1666 -- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
1667 foldl :: (a -> b -> a) -> a -> IntMap b -> a
1668 foldl f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1669 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1670 | otherwise -> go (go z l) r
1671 _ -> go z t
1672 where
1673 go z' Nil = z'
1674 go z' (Tip _ x) = f z' x
1675 go z' (Bin _ _ l r) = go (go z' l) r
1676 {-# INLINE foldl #-}
1677
1678 -- | /O(n)/. A strict version of 'foldl'. Each application of the operator is
1679 -- evaluated before using the result in the next application. This
1680 -- function is strict in the starting value.
1681 foldl' :: (a -> b -> a) -> a -> IntMap b -> a
1682 foldl' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1683 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1684 | otherwise -> go (go z l) r
1685 _ -> go z t
1686 where
1687 go !z' Nil = z'
1688 go z' (Tip _ x) = f z' x
1689 go z' (Bin _ _ l r) = go (go z' l) r
1690 {-# INLINE foldl' #-}
1691
1692 -- | /O(n)/. Fold the keys and values in the map using the given right-associative
1693 -- binary operator, such that
1694 -- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
1695 --
1696 -- For example,
1697 --
1698 -- > keys map = foldrWithKey (\k x ks -> k:ks) [] map
1699 --
1700 -- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
1701 -- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)"
1702 foldrWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
1703 foldrWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1704 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1705 | otherwise -> go (go z r) l
1706 _ -> go z t
1707 where
1708 go z' Nil = z'
1709 go z' (Tip kx x) = f kx x z'
1710 go z' (Bin _ _ l r) = go (go z' r) l
1711 {-# INLINE foldrWithKey #-}
1712
1713 -- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is
1714 -- evaluated before using the result in the next application. This
1715 -- function is strict in the starting value.
1716 foldrWithKey' :: (Key -> a -> b -> b) -> b -> IntMap a -> b
1717 foldrWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1718 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1719 | otherwise -> go (go z r) l
1720 _ -> go z t
1721 where
1722 go !z' Nil = z'
1723 go z' (Tip kx x) = f kx x z'
1724 go z' (Bin _ _ l r) = go (go z' r) l
1725 {-# INLINE foldrWithKey' #-}
1726
1727 -- | /O(n)/. Fold the keys and values in the map using the given left-associative
1728 -- binary operator, such that
1729 -- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@.
1730 --
1731 -- For example,
1732 --
1733 -- > keys = reverse . foldlWithKey (\ks k x -> k:ks) []
1734 --
1735 -- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
1736 -- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)"
1737 foldlWithKey :: (a -> Key -> b -> a) -> a -> IntMap b -> a
1738 foldlWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1739 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1740 | otherwise -> go (go z l) r
1741 _ -> go z t
1742 where
1743 go z' Nil = z'
1744 go z' (Tip kx x) = f z' kx x
1745 go z' (Bin _ _ l r) = go (go z' l) r
1746 {-# INLINE foldlWithKey #-}
1747
1748 -- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is
1749 -- evaluated before using the result in the next application. This
1750 -- function is strict in the starting value.
1751 foldlWithKey' :: (a -> Key -> b -> a) -> a -> IntMap b -> a
1752 foldlWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1753 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1754 | otherwise -> go (go z l) r
1755 _ -> go z t
1756 where
1757 go !z' Nil = z'
1758 go z' (Tip kx x) = f z' kx x
1759 go z' (Bin _ _ l r) = go (go z' l) r
1760 {-# INLINE foldlWithKey' #-}
1761
1762 -- | /O(n)/. Fold the keys and values in the map using the given monoid, such that
1763 --
1764 -- @'foldMapWithKey' f = 'Prelude.fold' . 'mapWithKey' f@
1765 --
1766 -- This can be an asymptotically faster than 'foldrWithKey' or 'foldlWithKey' for some monoids.
1767 foldMapWithKey :: Monoid m => (Key -> a -> m) -> IntMap a -> m
1768 foldMapWithKey f = go
1769 where
1770 go Nil = mempty
1771 go (Tip kx x) = f kx x
1772 go (Bin _ _ l r) = go l `mappend` go r
1773 {-# INLINE foldMapWithKey #-}
1774
1775 {--------------------------------------------------------------------
1776 List variations
1777 --------------------------------------------------------------------}
1778 -- | /O(n)/.
1779 -- Return all elements of the map in the ascending order of their keys.
1780 -- Subject to list fusion.
1781 --
1782 -- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"]
1783 -- > elems empty == []
1784
1785 elems :: IntMap a -> [a]
1786 elems = foldr (:) []
1787
1788 -- | /O(n)/. Return all keys of the map in ascending order. Subject to list
1789 -- fusion.
1790 --
1791 -- > keys (fromList [(5,"a"), (3,"b")]) == [3,5]
1792 -- > keys empty == []
1793
1794 keys :: IntMap a -> [Key]
1795 keys = foldrWithKey (\k _ ks -> k : ks) []
1796
1797 -- | /O(n)/. An alias for 'toAscList'. Returns all key\/value pairs in the
1798 -- map in ascending key order. Subject to list fusion.
1799 --
1800 -- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1801 -- > assocs empty == []
1802
1803 assocs :: IntMap a -> [(Key,a)]
1804 assocs = toAscList
1805
1806 -- | /O(n*min(n,W))/. The set of all keys of the map.
1807 --
1808 -- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.IntSet.fromList [3,5]
1809 -- > keysSet empty == Data.IntSet.empty
1810
1811 keysSet :: IntMap a -> IntSet.IntSet
1812 keysSet Nil = IntSet.Nil
1813 keysSet (Tip kx _) = IntSet.singleton kx
1814 keysSet (Bin p m l r)
1815 | m .&. IntSet.suffixBitMask == 0 = IntSet.Bin p m (keysSet l) (keysSet r)
1816 | otherwise = IntSet.Tip (p .&. IntSet.prefixBitMask) (computeBm (computeBm 0 l) r)
1817 where computeBm !acc (Bin _ _ l' r') = computeBm (computeBm acc l') r'
1818 computeBm acc (Tip kx _) = acc .|. IntSet.bitmapOf kx
1819 computeBm _ Nil = error "Data.IntSet.keysSet: Nil"
1820
1821 -- | /O(n)/. Build a map from a set of keys and a function which for each key
1822 -- computes its value.
1823 --
1824 -- > fromSet (\k -> replicate k 'a') (Data.IntSet.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")]
1825 -- > fromSet undefined Data.IntSet.empty == empty
1826
1827 fromSet :: (Key -> a) -> IntSet.IntSet -> IntMap a
1828 fromSet _ IntSet.Nil = Nil
1829 fromSet f (IntSet.Bin p m l r) = Bin p m (fromSet f l) (fromSet f r)
1830 fromSet f (IntSet.Tip kx bm) = buildTree f kx bm (IntSet.suffixBitMask + 1)
1831 where -- This is slightly complicated, as we to convert the dense
1832 -- representation of IntSet into tree representation of IntMap.
1833 --
1834 -- We are given a nonzero bit mask 'bmask' of 'bits' bits with prefix 'prefix'.
1835 -- We split bmask into halves corresponding to left and right subtree.
1836 -- If they are both nonempty, we create a Bin node, otherwise exactly
1837 -- one of them is nonempty and we construct the IntMap from that half.
1838 buildTree g !prefix !bmask bits = case bits of
1839 0 -> Tip prefix (g prefix)
1840 _ -> case intFromNat ((natFromInt bits) `shiftRL` 1) of
1841 bits2 | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 ->
1842 buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2
1843 | (bmask `shiftRL` bits2) .&. ((1 `shiftLL` bits2) - 1) == 0 ->
1844 buildTree g prefix bmask bits2
1845 | otherwise ->
1846 Bin prefix bits2 (buildTree g prefix bmask bits2) (buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2)
1847
1848 {--------------------------------------------------------------------
1849 Lists
1850 --------------------------------------------------------------------}
1851 #if __GLASGOW_HASKELL__ >= 708
1852 instance GHCExts.IsList (IntMap a) where
1853 type Item (IntMap a) = (Key,a)
1854 fromList = fromList
1855 toList = toList
1856 #endif
1857
1858 -- | /O(n)/. Convert the map to a list of key\/value pairs. Subject to list
1859 -- fusion.
1860 --
1861 -- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1862 -- > toList empty == []
1863
1864 toList :: IntMap a -> [(Key,a)]
1865 toList = toAscList
1866
1867 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
1868 -- keys are in ascending order. Subject to list fusion.
1869 --
1870 -- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1871
1872 toAscList :: IntMap a -> [(Key,a)]
1873 toAscList = foldrWithKey (\k x xs -> (k,x):xs) []
1874
1875 -- | /O(n)/. Convert the map to a list of key\/value pairs where the keys
1876 -- are in descending order. Subject to list fusion.
1877 --
1878 -- > toDescList (fromList [(5,"a"), (3,"b")]) == [(5,"a"), (3,"b")]
1879
1880 toDescList :: IntMap a -> [(Key,a)]
1881 toDescList = foldlWithKey (\xs k x -> (k,x):xs) []
1882
1883 -- List fusion for the list generating functions.
1884 #if __GLASGOW_HASKELL__
1885 -- The foldrFB and foldlFB are fold{r,l}WithKey equivalents, used for list fusion.
1886 -- They are important to convert unfused methods back, see mapFB in prelude.
1887 foldrFB :: (Key -> a -> b -> b) -> b -> IntMap a -> b
1888 foldrFB = foldrWithKey
1889 {-# INLINE[0] foldrFB #-}
1890 foldlFB :: (a -> Key -> b -> a) -> a -> IntMap b -> a
1891 foldlFB = foldlWithKey
1892 {-# INLINE[0] foldlFB #-}
1893
1894 -- Inline assocs and toList, so that we need to fuse only toAscList.
1895 {-# INLINE assocs #-}
1896 {-# INLINE toList #-}
1897
1898 -- The fusion is enabled up to phase 2 included. If it does not succeed,
1899 -- convert in phase 1 the expanded elems,keys,to{Asc,Desc}List calls back to
1900 -- elems,keys,to{Asc,Desc}List. In phase 0, we inline fold{lr}FB (which were
1901 -- used in a list fusion, otherwise it would go away in phase 1), and let compiler
1902 -- do whatever it wants with elems,keys,to{Asc,Desc}List -- it was forbidden to
1903 -- inline it before phase 0, otherwise the fusion rules would not fire at all.
1904 {-# NOINLINE[0] elems #-}
1905 {-# NOINLINE[0] keys #-}
1906 {-# NOINLINE[0] toAscList #-}
1907 {-# NOINLINE[0] toDescList #-}
1908 {-# RULES "IntMap.elems" [~1] forall m . elems m = build (\c n -> foldrFB (\_ x xs -> c x xs) n m) #-}
1909 {-# RULES "IntMap.elemsBack" [1] foldrFB (\_ x xs -> x : xs) [] = elems #-}
1910 {-# RULES "IntMap.keys" [~1] forall m . keys m = build (\c n -> foldrFB (\k _ xs -> c k xs) n m) #-}
1911 {-# RULES "IntMap.keysBack" [1] foldrFB (\k _ xs -> k : xs) [] = keys #-}
1912 {-# RULES "IntMap.toAscList" [~1] forall m . toAscList m = build (\c n -> foldrFB (\k x xs -> c (k,x) xs) n m) #-}
1913 {-# RULES "IntMap.toAscListBack" [1] foldrFB (\k x xs -> (k, x) : xs) [] = toAscList #-}
1914 {-# RULES "IntMap.toDescList" [~1] forall m . toDescList m = build (\c n -> foldlFB (\xs k x -> c (k,x) xs) n m) #-}
1915 {-# RULES "IntMap.toDescListBack" [1] foldlFB (\xs k x -> (k, x) : xs) [] = toDescList #-}
1916 #endif
1917
1918
1919 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
1920 --
1921 -- > fromList [] == empty
1922 -- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
1923 -- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
1924
1925 fromList :: [(Key,a)] -> IntMap a
1926 fromList xs
1927 = foldlStrict ins empty xs
1928 where
1929 ins t (k,x) = insert k x t
1930
1931 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
1932 --
1933 -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "ab"), (5, "cba")]
1934 -- > fromListWith (++) [] == empty
1935
1936 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1937 fromListWith f xs
1938 = fromListWithKey (\_ x y -> f x y) xs
1939
1940 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
1941 --
1942 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
1943 -- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "3:a|b"), (5, "5:c|5:b|a")]
1944 -- > fromListWithKey f [] == empty
1945
1946 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1947 fromListWithKey f xs
1948 = foldlStrict ins empty xs
1949 where
1950 ins t (k,x) = insertWithKey f k x t
1951
1952 -- | /O(n)/. Build a map from a list of key\/value pairs where
1953 -- the keys are in ascending order.
1954 --
1955 -- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
1956 -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")]
1957
1958 fromAscList :: [(Key,a)] -> IntMap a
1959 fromAscList xs
1960 = fromAscListWithKey (\_ x _ -> x) xs
1961
1962 -- | /O(n)/. Build a map from a list of key\/value pairs where
1963 -- the keys are in ascending order, with a combining function on equal keys.
1964 -- /The precondition (input list is ascending) is not checked./
1965 --
1966 -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
1967
1968 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1969 fromAscListWith f xs
1970 = fromAscListWithKey (\_ x y -> f x y) 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 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
1977 -- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "5:b|a")]
1978
1979 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1980 fromAscListWithKey _ [] = Nil
1981 fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0)
1982 where
1983 -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
1984 combineEq z [] = [z]
1985 combineEq z@(kz,zz) (x@(kx,xx):xs)
1986 | kx==kz = let yy = f kx xx zz in combineEq (kx,yy) xs
1987 | otherwise = z:combineEq x xs
1988
1989 -- | /O(n)/. Build a map from a list of key\/value pairs where
1990 -- the keys are in ascending order and all distinct.
1991 -- /The precondition (input list is strictly ascending) is not checked./
1992 --
1993 -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
1994
1995 fromDistinctAscList :: forall a. [(Key,a)] -> IntMap a
1996 fromDistinctAscList [] = Nil
1997 fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada
1998 where
1999 work (kx,vx) [] stk = finish kx (Tip kx vx) stk
2000 work (kx,vx) (z@(kz,_):zs) stk = reduce z zs (branchMask kx kz) kx (Tip kx vx) stk
2001
2002 reduce :: (Key,a) -> [(Key,a)] -> Mask -> Prefix -> IntMap a -> Stack a -> IntMap a
2003 reduce z zs _ px tx Nada = work z zs (Push px tx Nada)
2004 reduce z zs m px tx stk@(Push py ty stk') =
2005 let mxy = branchMask px py
2006 pxy = mask px mxy
2007 in if shorter m mxy
2008 then reduce z zs m pxy (Bin pxy mxy ty tx) stk'
2009 else work z zs (Push px tx stk)
2010
2011 finish _ t Nada = t
2012 finish px tx (Push py ty stk) = finish p (link py ty px tx) stk
2013 where m = branchMask px py
2014 p = mask px m
2015
2016 data Stack a = Push {-# UNPACK #-} !Prefix !(IntMap a) !(Stack a) | Nada
2017
2018
2019 {--------------------------------------------------------------------
2020 Eq
2021 --------------------------------------------------------------------}
2022 instance Eq a => Eq (IntMap a) where
2023 t1 == t2 = equal t1 t2
2024 t1 /= t2 = nequal t1 t2
2025
2026 equal :: Eq a => IntMap a -> IntMap a -> Bool
2027 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
2028 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
2029 equal (Tip kx x) (Tip ky y)
2030 = (kx == ky) && (x==y)
2031 equal Nil Nil = True
2032 equal _ _ = False
2033
2034 nequal :: Eq a => IntMap a -> IntMap a -> Bool
2035 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
2036 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
2037 nequal (Tip kx x) (Tip ky y)
2038 = (kx /= ky) || (x/=y)
2039 nequal Nil Nil = False
2040 nequal _ _ = True
2041
2042 {--------------------------------------------------------------------
2043 Ord
2044 --------------------------------------------------------------------}
2045
2046 instance Ord a => Ord (IntMap a) where
2047 compare m1 m2 = compare (toList m1) (toList m2)
2048
2049 {--------------------------------------------------------------------
2050 Functor
2051 --------------------------------------------------------------------}
2052
2053 instance Functor IntMap where
2054 fmap = map
2055
2056 {--------------------------------------------------------------------
2057 Show
2058 --------------------------------------------------------------------}
2059
2060 instance Show a => Show (IntMap a) where
2061 showsPrec d m = showParen (d > 10) $
2062 showString "fromList " . shows (toList m)
2063
2064 {--------------------------------------------------------------------
2065 Read
2066 --------------------------------------------------------------------}
2067 instance (Read e) => Read (IntMap e) where
2068 #ifdef __GLASGOW_HASKELL__
2069 readPrec = parens $ prec 10 $ do
2070 Ident "fromList" <- lexP
2071 xs <- readPrec
2072 return (fromList xs)
2073
2074 readListPrec = readListPrecDefault
2075 #else
2076 readsPrec p = readParen (p > 10) $ \ r -> do
2077 ("fromList",s) <- lex r
2078 (xs,t) <- reads s
2079 return (fromList xs,t)
2080 #endif
2081
2082 {--------------------------------------------------------------------
2083 Typeable
2084 --------------------------------------------------------------------}
2085
2086 INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap")
2087
2088 {--------------------------------------------------------------------
2089 Helpers
2090 --------------------------------------------------------------------}
2091 {--------------------------------------------------------------------
2092 Link
2093 --------------------------------------------------------------------}
2094 link :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
2095 link p1 t1 p2 t2
2096 | zero p1 m = Bin p m t1 t2
2097 | otherwise = Bin p m t2 t1
2098 where
2099 m = branchMask p1 p2
2100 p = mask p1 m
2101 {-# INLINE link #-}
2102
2103 {--------------------------------------------------------------------
2104 @bin@ assures that we never have empty trees within a tree.
2105 --------------------------------------------------------------------}
2106 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
2107 bin _ _ l Nil = l
2108 bin _ _ Nil r = r
2109 bin p m l r = Bin p m l r
2110 {-# INLINE bin #-}
2111
2112 -- binCheckLeft only checks that the left subtree is non-empty
2113 binCheckLeft :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
2114 binCheckLeft _ _ Nil r = r
2115 binCheckLeft p m l r = Bin p m l r
2116 {-# INLINE binCheckLeft #-}
2117
2118 -- binCheckRight only checks that the right subtree is non-empty
2119 binCheckRight :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
2120 binCheckRight _ _ l Nil = l
2121 binCheckRight p m l r = Bin p m l r
2122 {-# INLINE binCheckRight #-}
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