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