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