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