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