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