Fix forgotten documentation of Int{Map,Set}.delete{Min,Max}.
[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 -- ** Strict folds
127 , foldr'
128 , foldl'
129 , foldrWithKey'
130 , foldlWithKey'
131
132 -- * Conversion
133 , elems
134 , keys
135 , assocs
136 , keysSet
137 , fromSet
138
139 -- ** Lists
140 , toList
141 , fromList
142 , fromListWith
143 , fromListWithKey
144
145 -- ** Ordered lists
146 , toAscList
147 , toDescList
148 , fromAscList
149 , fromAscListWith
150 , fromAscListWithKey
151 , fromDistinctAscList
152
153 -- * Filter
154 , filter
155 , filterWithKey
156 , partition
157 , partitionWithKey
158
159 , mapMaybe
160 , mapMaybeWithKey
161 , mapEither
162 , mapEitherWithKey
163
164 , split
165 , splitLookup
166
167 -- * Submap
168 , isSubmapOf, isSubmapOfBy
169 , isProperSubmapOf, isProperSubmapOfBy
170
171 -- * Min\/Max
172 , findMin
173 , findMax
174 , deleteMin
175 , deleteMax
176 , deleteFindMin
177 , deleteFindMax
178 , updateMin
179 , updateMax
180 , updateMinWithKey
181 , updateMaxWithKey
182 , minView
183 , maxView
184 , minViewWithKey
185 , maxViewWithKey
186
187 -- * Debugging
188 , showTree
189 , showTreeWith
190
191 -- * Internal types
192 , Mask, Prefix, Nat
193
194 -- * Utility
195 , natFromInt
196 , intFromNat
197 , shiftRL
198 , shiftLL
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 Data.Bits
213
214 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
215 import qualified Data.IntSet.Base as IntSet
216 import Data.Monoid (Monoid(..))
217 import Data.Maybe (fromMaybe)
218 import Data.Typeable
219 import qualified Data.Foldable as Foldable
220 import Data.Traversable (Traversable(traverse))
221 import Control.Applicative (Applicative(pure,(<*>)),(<$>))
222 import Control.Monad ( liftM )
223 import Control.DeepSeq (NFData(rnf))
224
225 import Data.IntSet.Base (Key)
226 import Data.StrictPair
227
228 #if __GLASGOW_HASKELL__
229 import Text.Read
230 import Data.Data (Data(..), mkNoRepType)
231 #endif
232
233 #if __GLASGOW_HASKELL__
234 import GHC.Exts ( Word(..), Int(..), build )
235 import GHC.Prim ( uncheckedShiftL#, uncheckedShiftRL# )
236 #else
237 import Data.Word
238 #endif
239
240 -- On GHC, include MachDeps.h to get WORD_SIZE_IN_BITS macro.
241 #if defined(__GLASGOW_HASKELL__)
242 #include "MachDeps.h"
243 #endif
244
245 -- Use macros to define strictness of functions.
246 -- STRICT_x_OF_y denotes an y-ary function strict in the x-th parameter.
247 -- We do not use BangPatterns, because they are not in any standard and we
248 -- want the compilers to be compiled by as many compilers as possible.
249 #define STRICT_1_OF_2(fn) fn arg _ | arg `seq` False = undefined
250
251 -- A "Nat" is a natural machine word (an unsigned Int)
252 type Nat = Word
253
254 natFromInt :: Key -> Nat
255 natFromInt = fromIntegral
256 {-# INLINE natFromInt #-}
257
258 intFromNat :: Nat -> Key
259 intFromNat = fromIntegral
260 {-# INLINE intFromNat #-}
261
262 -- Right and left logical shifts.
263 shiftRL, shiftLL :: Nat -> Key -> Nat
264 #if __GLASGOW_HASKELL__
265 {--------------------------------------------------------------------
266 GHC: use unboxing to get @shiftRL@ inlined.
267 --------------------------------------------------------------------}
268 shiftRL (W# x) (I# i) = W# (uncheckedShiftRL# x i)
269 shiftLL (W# x) (I# i) = W# (uncheckedShiftL# x i)
270 #else
271 shiftRL x i = shiftR x i
272 shiftLL x i = shiftL x i
273 #endif
274 {-# INLINE shiftRL #-}
275 {-# INLINE shiftLL #-}
276
277 {--------------------------------------------------------------------
278 Types
279 --------------------------------------------------------------------}
280
281
282 -- | A map of integers to values @a@.
283
284 -- See Note: Order of constructors
285 data IntMap a = Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(IntMap a) !(IntMap a)
286 | Tip {-# UNPACK #-} !Key a
287 | Nil
288
289 type Prefix = Int
290 type Mask = Int
291
292 {--------------------------------------------------------------------
293 Operators
294 --------------------------------------------------------------------}
295
296 -- | /O(min(n,W))/. Find the value at a key.
297 -- Calls 'error' when the element can not be found.
298 --
299 -- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map
300 -- > fromList [(5,'a'), (3,'b')] ! 5 == 'a'
301
302 (!) :: IntMap a -> Key -> a
303 m ! k = find k m
304
305 -- | Same as 'difference'.
306 (\\) :: IntMap a -> IntMap b -> IntMap a
307 m1 \\ m2 = difference m1 m2
308
309 infixl 9 \\{-This comment teaches CPP correct behaviour -}
310
311 {--------------------------------------------------------------------
312 Types
313 --------------------------------------------------------------------}
314
315 instance Monoid (IntMap a) where
316 mempty = empty
317 mappend = union
318 mconcat = unions
319
320 instance Foldable.Foldable IntMap where
321 fold Nil = mempty
322 fold (Tip _ v) = v
323 fold (Bin _ _ l r) = Foldable.fold l `mappend` Foldable.fold r
324 foldr = foldr
325 foldl = foldl
326 foldMap _ Nil = mempty
327 foldMap f (Tip _k v) = f v
328 foldMap f (Bin _ _ l r) = Foldable.foldMap f l `mappend` Foldable.foldMap f r
329
330 instance Traversable IntMap where
331 traverse f = traverseWithKey (\_ -> f)
332
333 instance NFData a => NFData (IntMap a) where
334 rnf Nil = ()
335 rnf (Tip _ v) = rnf v
336 rnf (Bin _ _ l r) = rnf l `seq` rnf r
337
338 #if __GLASGOW_HASKELL__
339
340 {--------------------------------------------------------------------
341 A Data instance
342 --------------------------------------------------------------------}
343
344 -- This instance preserves data abstraction at the cost of inefficiency.
345 -- We omit reflection services for the sake of data abstraction.
346
347 instance Data a => Data (IntMap a) where
348 gfoldl f z im = z fromList `f` (toList im)
349 toConstr _ = error "toConstr"
350 gunfold _ _ = error "gunfold"
351 dataTypeOf _ = mkNoRepType "Data.IntMap.IntMap"
352 dataCast1 f = gcast1 f
353
354 #endif
355
356 {--------------------------------------------------------------------
357 Query
358 --------------------------------------------------------------------}
359 -- | /O(1)/. Is the map empty?
360 --
361 -- > Data.IntMap.null (empty) == True
362 -- > Data.IntMap.null (singleton 1 'a') == False
363
364 null :: IntMap a -> Bool
365 null Nil = True
366 null _ = False
367 {-# INLINE null #-}
368
369 -- | /O(n)/. Number of elements in the map.
370 --
371 -- > size empty == 0
372 -- > size (singleton 1 'a') == 1
373 -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3
374 size :: IntMap a -> Int
375 size t
376 = case t of
377 Bin _ _ l r -> size l + size r
378 Tip _ _ -> 1
379 Nil -> 0
380
381 -- | /O(min(n,W))/. Is the key a member of the map?
382 --
383 -- > member 5 (fromList [(5,'a'), (3,'b')]) == True
384 -- > member 1 (fromList [(5,'a'), (3,'b')]) == False
385
386 -- See Note: Local 'go' functions and capturing]
387 member :: Key -> IntMap a -> Bool
388 member k = k `seq` go
389 where
390 go (Bin p m l r) | nomatch k p m = False
391 | zero k m = go l
392 | otherwise = go r
393 go (Tip kx _) = k == kx
394 go Nil = False
395
396 -- | /O(min(n,W))/. Is the key not a member of the map?
397 --
398 -- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False
399 -- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True
400
401 notMember :: Key -> IntMap a -> Bool
402 notMember k m = not $ member k m
403
404 -- | /O(min(n,W))/. Lookup the value at a key in the map. See also 'Data.Map.lookup'.
405
406 -- See Note: Local 'go' functions and capturing]
407 lookup :: Key -> IntMap a -> Maybe a
408 lookup k = k `seq` go
409 where
410 go (Bin p m l r) | nomatch k p m = Nothing
411 | zero k m = go l
412 | otherwise = go r
413 go (Tip kx x) | k == kx = Just x
414 | otherwise = Nothing
415 go Nil = Nothing
416
417
418 -- See Note: Local 'go' functions and capturing]
419 find :: Key -> IntMap a -> a
420 find k = k `seq` go
421 where
422 go (Bin p m l r) | nomatch k p m = not_found
423 | zero k m = go l
424 | otherwise = go r
425 go (Tip kx x) | k == kx = x
426 | otherwise = not_found
427 go Nil = not_found
428
429 not_found = error ("IntMap.!: key " ++ show k ++ " is not an element of the map")
430
431 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
432 -- returns the value at key @k@ or returns @def@ when the key is not an
433 -- element of the map.
434 --
435 -- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x'
436 -- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a'
437
438 -- See Note: Local 'go' functions and capturing]
439 findWithDefault :: a -> Key -> IntMap a -> a
440 findWithDefault def k = k `seq` go
441 where
442 go (Bin p m l r) | nomatch k p m = def
443 | zero k m = go l
444 | otherwise = go r
445 go (Tip kx x) | k == kx = x
446 | otherwise = def
447 go Nil = def
448
449 -- | /O(log n)/. Find largest key smaller than the given one and return the
450 -- corresponding (key, value) pair.
451 --
452 -- > lookupLT 3 (fromList [(3,'a'), (5,'b')]) == Nothing
453 -- > lookupLT 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
454
455 -- See Note: Local 'go' functions and capturing.
456 lookupLT :: Key -> IntMap a -> Maybe (Key, a)
457 lookupLT k t = k `seq` case t of
458 Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r
459 _ -> go Nil t
460 where
461 go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMax def else unsafeFindMax r
462 | zero k m = go def l
463 | otherwise = go l r
464 go def (Tip ky y) | k <= ky = unsafeFindMax def
465 | otherwise = Just (ky, y)
466 go def Nil = unsafeFindMax def
467
468 -- | /O(log n)/. Find smallest key greater than the given one and return the
469 -- corresponding (key, value) pair.
470 --
471 -- > lookupGT 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
472 -- > lookupGT 5 (fromList [(3,'a'), (5,'b')]) == Nothing
473
474 -- See Note: Local 'go' functions and capturing.
475 lookupGT :: Key -> IntMap a -> Maybe (Key, a)
476 lookupGT k t = k `seq` case t of
477 Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r
478 _ -> go Nil t
479 where
480 go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMin l else unsafeFindMin def
481 | zero k m = go r l
482 | otherwise = go def r
483 go def (Tip ky y) | k >= ky = unsafeFindMin def
484 | otherwise = Just (ky, y)
485 go def Nil = unsafeFindMin def
486
487 -- | /O(log n)/. Find largest key smaller or equal to the given one and return
488 -- the corresponding (key, value) pair.
489 --
490 -- > lookupLE 2 (fromList [(3,'a'), (5,'b')]) == Nothing
491 -- > lookupLE 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
492 -- > lookupLE 5 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
493
494 -- See Note: Local 'go' functions and capturing.
495 lookupLE :: Key -> IntMap a -> Maybe (Key, a)
496 lookupLE k t = k `seq` case t of
497 Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r
498 _ -> go Nil t
499 where
500 go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMax def else unsafeFindMax r
501 | zero k m = go def l
502 | otherwise = go l r
503 go def (Tip ky y) | k < ky = unsafeFindMax def
504 | otherwise = Just (ky, y)
505 go def Nil = unsafeFindMax def
506
507 -- | /O(log n)/. Find smallest key greater or equal to the given one and return
508 -- the corresponding (key, value) pair.
509 --
510 -- > lookupGE 3 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
511 -- > lookupGE 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
512 -- > lookupGE 6 (fromList [(3,'a'), (5,'b')]) == Nothing
513
514 -- See Note: Local 'go' functions and capturing.
515 lookupGE :: Key -> IntMap a -> Maybe (Key, a)
516 lookupGE k t = k `seq` case t of
517 Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r
518 _ -> go Nil t
519 where
520 go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMin l else unsafeFindMin def
521 | zero k m = go r l
522 | otherwise = go def r
523 go def (Tip ky y) | k > ky = unsafeFindMin def
524 | otherwise = Just (ky, y)
525 go def Nil = unsafeFindMin def
526
527
528 -- Helper function for lookupGE and lookupGT. It assumes that if a Bin node is
529 -- given, it has m > 0.
530 unsafeFindMin :: IntMap a -> Maybe (Key, a)
531 unsafeFindMin Nil = Nothing
532 unsafeFindMin (Tip ky y) = Just (ky, y)
533 unsafeFindMin (Bin _ _ l _) = unsafeFindMin l
534
535 -- Helper function for lookupLE and lookupLT. It assumes that if a Bin node is
536 -- given, it has m > 0.
537 unsafeFindMax :: IntMap a -> Maybe (Key, a)
538 unsafeFindMax Nil = Nothing
539 unsafeFindMax (Tip ky y) = Just (ky, y)
540 unsafeFindMax (Bin _ _ _ r) = unsafeFindMax r
541
542 {--------------------------------------------------------------------
543 Construction
544 --------------------------------------------------------------------}
545 -- | /O(1)/. The empty map.
546 --
547 -- > empty == fromList []
548 -- > size empty == 0
549
550 empty :: IntMap a
551 empty
552 = Nil
553 {-# INLINE empty #-}
554
555 -- | /O(1)/. A map of one element.
556 --
557 -- > singleton 1 'a' == fromList [(1, 'a')]
558 -- > size (singleton 1 'a') == 1
559
560 singleton :: Key -> a -> IntMap a
561 singleton k x
562 = Tip k x
563 {-# INLINE singleton #-}
564
565 {--------------------------------------------------------------------
566 Insert
567 --------------------------------------------------------------------}
568 -- | /O(min(n,W))/. Insert a new key\/value pair in the map.
569 -- If the key is already present in the map, the associated value is
570 -- replaced with the supplied value, i.e. 'insert' is equivalent to
571 -- @'insertWith' 'const'@.
572 --
573 -- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')]
574 -- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')]
575 -- > insert 5 'x' empty == singleton 5 'x'
576
577 insert :: Key -> a -> IntMap a -> IntMap a
578 insert k x t = k `seq`
579 case t of
580 Bin p m l r
581 | nomatch k p m -> join k (Tip k x) p t
582 | zero k m -> Bin p m (insert k x l) r
583 | otherwise -> Bin p m l (insert k x r)
584 Tip ky _
585 | k==ky -> Tip k x
586 | otherwise -> join k (Tip k x) ky t
587 Nil -> Tip k x
588
589 -- right-biased insertion, used by 'union'
590 -- | /O(min(n,W))/. Insert with a combining function.
591 -- @'insertWith' f key value mp@
592 -- will insert the pair (key, value) into @mp@ if key does
593 -- not exist in the map. If the key does exist, the function will
594 -- insert @f new_value old_value@.
595 --
596 -- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")]
597 -- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
598 -- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx"
599
600 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
601 insertWith f k x t
602 = insertWithKey (\_ x' y' -> f x' y') k x t
603
604 -- | /O(min(n,W))/. Insert with a combining function.
605 -- @'insertWithKey' f key value mp@
606 -- will insert the pair (key, value) into @mp@ if key does
607 -- not exist in the map. If the key does exist, the function will
608 -- insert @f key new_value old_value@.
609 --
610 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
611 -- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")]
612 -- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
613 -- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx"
614
615 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
616 insertWithKey f k x t = k `seq`
617 case t of
618 Bin p m l r
619 | nomatch k p m -> join k (Tip k x) p t
620 | zero k m -> Bin p m (insertWithKey f k x l) r
621 | otherwise -> Bin p m l (insertWithKey f k x r)
622 Tip ky y
623 | k==ky -> Tip k (f k x y)
624 | otherwise -> join k (Tip k x) ky t
625 Nil -> Tip k x
626
627 -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
628 -- is a pair where the first element is equal to (@'lookup' k map@)
629 -- and the second element equal to (@'insertWithKey' f k x map@).
630 --
631 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
632 -- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")])
633 -- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")])
634 -- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx")
635 --
636 -- This is how to define @insertLookup@ using @insertLookupWithKey@:
637 --
638 -- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t
639 -- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")])
640 -- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")])
641
642 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
643 insertLookupWithKey f k x t = k `seq`
644 case t of
645 Bin p m l r
646 | nomatch k p m -> (Nothing,join k (Tip k x) p t)
647 | zero k m -> let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r)
648 | otherwise -> let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r')
649 Tip ky y
650 | k==ky -> (Just y,Tip k (f k x y))
651 | otherwise -> (Nothing,join k (Tip k x) ky t)
652 Nil -> (Nothing,Tip k x)
653
654
655 {--------------------------------------------------------------------
656 Deletion
657 --------------------------------------------------------------------}
658 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
659 -- a member of the map, the original map is returned.
660 --
661 -- > delete 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
662 -- > delete 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
663 -- > delete 5 empty == empty
664
665 delete :: Key -> IntMap a -> IntMap a
666 delete k t = k `seq`
667 case t of
668 Bin p m l r
669 | nomatch k p m -> t
670 | zero k m -> bin p m (delete k l) r
671 | otherwise -> bin p m l (delete k r)
672 Tip ky _
673 | k==ky -> Nil
674 | otherwise -> t
675 Nil -> Nil
676
677 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
678 -- a member of the map, the original map is returned.
679 --
680 -- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
681 -- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
682 -- > adjust ("new " ++) 7 empty == empty
683
684 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
685 adjust f k m
686 = adjustWithKey (\_ x -> f x) k m
687
688 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
689 -- a member of the map, the original map is returned.
690 --
691 -- > let f key x = (show key) ++ ":new " ++ x
692 -- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
693 -- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
694 -- > adjustWithKey f 7 empty == empty
695
696 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
697 adjustWithKey f
698 = updateWithKey (\k' x -> Just (f k' 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 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 x = if x == "a" then Just "new a" else Nothing
705 -- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
706 -- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
707 -- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
708
709 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
710 update f
711 = updateWithKey (\_ x -> f x)
712
713 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
714 -- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
715 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
716 --
717 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
718 -- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
719 -- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
720 -- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
721
722 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
723 updateWithKey f k t = k `seq`
724 case t of
725 Bin p m l r
726 | nomatch k p m -> t
727 | zero k m -> bin p m (updateWithKey f k l) r
728 | otherwise -> bin p m l (updateWithKey f k r)
729 Tip ky y
730 | k==ky -> case (f k y) of
731 Just y' -> Tip ky y'
732 Nothing -> Nil
733 | otherwise -> t
734 Nil -> Nil
735
736 -- | /O(min(n,W))/. Lookup and update.
737 -- The function returns original value, if it is updated.
738 -- This is different behavior than 'Data.Map.updateLookupWithKey'.
739 -- Returns the original key value if the map entry is deleted.
740 --
741 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
742 -- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:new a")])
743 -- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")])
744 -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a")
745
746 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
747 updateLookupWithKey f k t = k `seq`
748 case t of
749 Bin p m l r
750 | nomatch k p m -> (Nothing,t)
751 | zero k m -> let (found,l') = updateLookupWithKey f k l in (found,bin p m l' r)
752 | otherwise -> let (found,r') = updateLookupWithKey f k r in (found,bin p m l r')
753 Tip ky y
754 | k==ky -> case (f k y) of
755 Just y' -> (Just y,Tip ky y')
756 Nothing -> (Just y,Nil)
757 | otherwise -> (Nothing,t)
758 Nil -> (Nothing,Nil)
759
760
761
762 -- | /O(min(n,W))/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
763 -- 'alter' can be used to insert, delete, or update a value in an 'IntMap'.
764 -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
765 alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
766 alter f k t = k `seq`
767 case t of
768 Bin p m l r
769 | nomatch k p m -> case f Nothing of
770 Nothing -> t
771 Just x -> join k (Tip k x) p t
772 | zero k m -> bin p m (alter f k l) r
773 | otherwise -> bin p m l (alter f k r)
774 Tip ky y
775 | k==ky -> case f (Just y) of
776 Just x -> Tip ky x
777 Nothing -> Nil
778 | otherwise -> case f Nothing of
779 Just x -> join k (Tip k x) ky t
780 Nothing -> Tip ky y
781 Nil -> case f Nothing of
782 Just x -> Tip k x
783 Nothing -> Nil
784
785
786 {--------------------------------------------------------------------
787 Union
788 --------------------------------------------------------------------}
789 -- | The union of a list of maps.
790 --
791 -- > unions [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
792 -- > == fromList [(3, "b"), (5, "a"), (7, "C")]
793 -- > unions [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])]
794 -- > == fromList [(3, "B3"), (5, "A3"), (7, "C")]
795
796 unions :: [IntMap a] -> IntMap a
797 unions xs
798 = foldlStrict union empty xs
799
800 -- | The union of a list of maps, with a combining operation.
801 --
802 -- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
803 -- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")]
804
805 unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a
806 unionsWith f ts
807 = foldlStrict (unionWith f) empty ts
808
809 -- | /O(n+m)/. The (left-biased) union of two maps.
810 -- It prefers the first map when duplicate keys are encountered,
811 -- i.e. (@'union' == 'unionWith' 'const'@).
812 --
813 -- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")]
814
815 union :: IntMap a -> IntMap a -> IntMap a
816 union m1 m2
817 = mergeWithKey' Bin const id id m1 m2
818
819 -- | /O(n+m)/. The union with a combining function.
820 --
821 -- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")]
822
823 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
824 unionWith f m1 m2
825 = unionWithKey (\_ x y -> f x y) m1 m2
826
827 -- | /O(n+m)/. The union with a combining function.
828 --
829 -- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value
830 -- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")]
831
832 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
833 unionWithKey f m1 m2
834 = mergeWithKey' Bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 (f k1 x1 x2)) id id m1 m2
835
836 {--------------------------------------------------------------------
837 Difference
838 --------------------------------------------------------------------}
839 -- | /O(n+m)/. Difference between two maps (based on keys).
840 --
841 -- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b"
842
843 difference :: IntMap a -> IntMap b -> IntMap a
844 difference m1 m2
845 = mergeWithKey (\_ _ _ -> Nothing) id (const Nil) m1 m2
846
847 -- | /O(n+m)/. Difference with a combining function.
848 --
849 -- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing
850 -- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")])
851 -- > == singleton 3 "b:B"
852
853 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
854 differenceWith f m1 m2
855 = differenceWithKey (\_ x y -> f x y) m1 m2
856
857 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
858 -- encountered, the combining function is applied to the key and both values.
859 -- If it returns 'Nothing', the element is discarded (proper set difference).
860 -- If it returns (@'Just' y@), the element is updated with a new value @y@.
861 --
862 -- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing
863 -- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")])
864 -- > == singleton 3 "3:b|B"
865
866 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
867 differenceWithKey f m1 m2
868 = mergeWithKey f id (const Nil) m1 m2
869
870
871 {--------------------------------------------------------------------
872 Intersection
873 --------------------------------------------------------------------}
874 -- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys).
875 --
876 -- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a"
877
878 intersection :: IntMap a -> IntMap b -> IntMap a
879 intersection m1 m2
880 = mergeWithKey' bin const (const Nil) (const Nil) m1 m2
881
882 -- | /O(n+m)/. The intersection with a combining function.
883 --
884 -- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA"
885
886 intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
887 intersectionWith f m1 m2
888 = intersectionWithKey (\_ x y -> f x y) m1 m2
889
890 -- | /O(n+m)/. The intersection with a combining function.
891 --
892 -- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar
893 -- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A"
894
895 intersectionWithKey :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
896 intersectionWithKey f m1 m2
897 = mergeWithKey' bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 (f k1 x1 x2)) (const Nil) (const Nil) m1 m2
898
899 {--------------------------------------------------------------------
900 MergeWithKey
901 --------------------------------------------------------------------}
902
903 -- | /O(n+m)/. A high-performance universal combining function. Using
904 -- 'mergeWithKey', all combining functions can be defined without any loss of
905 -- efficiency (with exception of 'union', 'difference' and 'intersection',
906 -- where sharing of some nodes is lost with 'mergeWithKey').
907 --
908 -- Please make sure you know what is going on when using 'mergeWithKey',
909 -- otherwise you can be surprised by unexpected code growth or even
910 -- corruption of the data structure.
911 --
912 -- When 'mergeWithKey' is given three arguments, it is inlined to the call
913 -- site. You should therefore use 'mergeWithKey' only to define your custom
914 -- combining functions. For example, you could define 'unionWithKey',
915 -- 'differenceWithKey' and 'intersectionWithKey' as
916 --
917 -- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2
918 -- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2
919 -- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2
920 --
921 -- When calling @'mergeWithKey' combine only1 only2@, a function combining two
922 -- 'IntMap's is created, such that
923 --
924 -- * if a key is present in both maps, it is passed with both corresponding
925 -- values to the @combine@ function. Depending on the result, the key is either
926 -- present in the result with specified value, or is left out;
927 --
928 -- * a nonempty subtree present only in the first map is passed to @only1@ and
929 -- the output is added to the result;
930 --
931 -- * a nonempty subtree present only in the second map is passed to @only2@ and
932 -- the output is added to the result.
933 --
934 -- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/.
935 -- The values can be modified arbitrarily. Most common variants of @only1@ and
936 -- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@ or
937 -- @'filterWithKey' f@ could be used for any @f@.
938
939 mergeWithKey :: (Key -> a -> b -> Maybe c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c)
940 -> IntMap a -> IntMap b -> IntMap c
941 mergeWithKey f g1 g2 = mergeWithKey' bin combine g1 g2
942 where -- We use the lambda form to avoid non-exhaustive pattern matches warning.
943 combine = \(Tip k1 x1) (Tip _k2 x2) -> case f k1 x1 x2 of Nothing -> Nil
944 Just x -> Tip k1 x
945 {-# INLINE combine #-}
946 {-# INLINE mergeWithKey #-}
947
948 -- Slightly more general version of mergeWithKey. It differs in the following:
949 --
950 -- * the combining function operates on maps instead of keys and values. The
951 -- reason is to enable sharing in union, difference and intersection.
952 --
953 -- * mergeWithKey' is given an equivalent of bin. The reason is that in union*,
954 -- Bin constructor can be used, because we know both subtrees are nonempty.
955
956 mergeWithKey' :: (Prefix -> Mask -> IntMap c -> IntMap c -> IntMap c)
957 -> (IntMap a -> IntMap b -> IntMap c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c)
958 -> IntMap a -> IntMap b -> IntMap c
959 mergeWithKey' bin' f g1 g2 = go
960 where
961 go t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
962 | shorter m1 m2 = merge1
963 | shorter m2 m1 = merge2
964 | p1 == p2 = bin' p1 m1 (go l1 l2) (go r1 r2)
965 | otherwise = maybe_join p1 (g1 t1) p2 (g2 t2)
966 where
967 merge1 | nomatch p2 p1 m1 = maybe_join p1 (g1 t1) p2 (g2 t2)
968 | zero p2 m1 = bin' p1 m1 (go l1 t2) (g1 r1)
969 | otherwise = bin' p1 m1 (g1 l1) (go r1 t2)
970 merge2 | nomatch p1 p2 m2 = maybe_join p1 (g1 t1) p2 (g2 t2)
971 | zero p1 m2 = bin' p2 m2 (go t1 l2) (g2 r2)
972 | otherwise = bin' p2 m2 (g2 l2) (go t1 r2)
973
974 go t1'@(Bin _ _ _ _) t2'@(Tip k2' _) = merge t2' k2' t1'
975 where merge t2 k2 t1@(Bin p1 m1 l1 r1) | nomatch k2 p1 m1 = maybe_join p1 (g1 t1) k2 (g2 t2)
976 | zero k2 m1 = bin' p1 m1 (merge t2 k2 l1) (g1 r1)
977 | otherwise = bin' p1 m1 (g1 l1) (merge t2 k2 r1)
978 merge t2 k2 t1@(Tip k1 _) | k1 == k2 = f t1 t2
979 | otherwise = maybe_join k1 (g1 t1) k2 (g2 t2)
980 merge t2 _ Nil = g2 t2
981
982 go t1@(Bin _ _ _ _) Nil = g1 t1
983
984 go t1'@(Tip k1' _) t2' = merge t1' k1' t2'
985 where merge t1 k1 t2@(Bin p2 m2 l2 r2) | nomatch k1 p2 m2 = maybe_join k1 (g1 t1) p2 (g2 t2)
986 | zero k1 m2 = bin' p2 m2 (merge t1 k1 l2) (g2 r2)
987 | otherwise = bin' p2 m2 (g2 l2) (merge t1 k1 r2)
988 merge t1 k1 t2@(Tip k2 _) | k1 == k2 = f t1 t2
989 | otherwise = maybe_join k1 (g1 t1) k2 (g2 t2)
990 merge t1 _ Nil = g1 t1
991
992 go Nil t2 = g2 t2
993
994 maybe_join _ Nil _ t2 = t2
995 maybe_join _ t1 _ Nil = t1
996 maybe_join p1 t1 p2 t2 = join p1 t1 p2 t2
997 {-# INLINE maybe_join #-}
998 {-# INLINE mergeWithKey' #-}
999
1000 {--------------------------------------------------------------------
1001 Min\/Max
1002 --------------------------------------------------------------------}
1003
1004 -- | /O(min(n,W))/. Update the value at the minimal key.
1005 --
1006 -- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"3:b"), (5,"a")]
1007 -- > updateMinWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1008
1009 updateMinWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a
1010 updateMinWithKey f t =
1011 case t of Bin p m l r | m < 0 -> bin p m l (go f r)
1012 _ -> go f t
1013 where
1014 go f' (Bin p m l r) = bin p m (go f' l) r
1015 go f' (Tip k y) = case f' k y of
1016 Just y' -> Tip k y'
1017 Nothing -> Nil
1018 go _ Nil = error "updateMinWithKey Nil"
1019
1020 -- | /O(min(n,W))/. Update the value at the maximal key.
1021 --
1022 -- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")]
1023 -- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1024
1025 updateMaxWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a
1026 updateMaxWithKey f t =
1027 case t of Bin p m l r | m < 0 -> bin p m (go f l) r
1028 _ -> go f t
1029 where
1030 go f' (Bin p m l r) = bin p m l (go f' r)
1031 go f' (Tip k y) = case f' k y of
1032 Just y' -> Tip k y'
1033 Nothing -> Nil
1034 go _ Nil = error "updateMaxWithKey Nil"
1035
1036 -- | /O(min(n,W))/. Retrieves the maximal (key,value) pair of the map, and
1037 -- the map stripped of that element, or 'Nothing' if passed an empty map.
1038 --
1039 -- > maxViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((5,"a"), singleton 3 "b")
1040 -- > maxViewWithKey empty == Nothing
1041
1042 maxViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
1043 maxViewWithKey t =
1044 case t of Nil -> Nothing
1045 Bin p m l r | m < 0 -> case go l of (result, l') -> Just (result, bin p m l' r)
1046 _ -> Just (go t)
1047 where
1048 go (Bin p m l r) = case go r of (result, r') -> (result, bin p m l r')
1049 go (Tip k y) = ((k, y), Nil)
1050 go Nil = error "maxViewWithKey Nil"
1051
1052 -- | /O(min(n,W))/. Retrieves the minimal (key,value) pair of the map, and
1053 -- the map stripped of that element, or 'Nothing' if passed an empty map.
1054 --
1055 -- > minViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((3,"b"), singleton 5 "a")
1056 -- > minViewWithKey empty == Nothing
1057
1058 minViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
1059 minViewWithKey t =
1060 case t of Nil -> Nothing
1061 Bin p m l r | m < 0 -> case go r of (result, r') -> Just (result, bin p m l r')
1062 _ -> Just (go t)
1063 where
1064 go (Bin p m l r) = case go l of (result, l') -> (result, bin p m l' r)
1065 go (Tip k y) = ((k, y), Nil)
1066 go Nil = error "minViewWithKey Nil"
1067
1068 -- | /O(min(n,W))/. Update the value at the maximal key.
1069 --
1070 -- > updateMax (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "Xa")]
1071 -- > updateMax (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1072
1073 updateMax :: (a -> Maybe a) -> IntMap a -> IntMap a
1074 updateMax f = updateMaxWithKey (const f)
1075
1076 -- | /O(min(n,W))/. Update the value at the minimal key.
1077 --
1078 -- > updateMin (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "Xb"), (5, "a")]
1079 -- > updateMin (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1080
1081 updateMin :: (a -> Maybe a) -> IntMap a -> IntMap a
1082 updateMin f = updateMinWithKey (const f)
1083
1084 -- Similar to the Arrow instance.
1085 first :: (a -> c) -> (a, b) -> (c, b)
1086 first f (x,y) = (f x,y)
1087
1088 -- | /O(min(n,W))/. Retrieves the maximal key of the map, and the map
1089 -- stripped of that element, or 'Nothing' if passed an empty map.
1090 maxView :: IntMap a -> Maybe (a, IntMap a)
1091 maxView t = liftM (first snd) (maxViewWithKey t)
1092
1093 -- | /O(min(n,W))/. Retrieves the minimal key of the map, and the map
1094 -- stripped of that element, or 'Nothing' if passed an empty map.
1095 minView :: IntMap a -> Maybe (a, IntMap a)
1096 minView t = liftM (first snd) (minViewWithKey t)
1097
1098 -- | /O(min(n,W))/. Delete and find the maximal element.
1099 deleteFindMax :: IntMap a -> ((Key, a), IntMap a)
1100 deleteFindMax = fromMaybe (error "deleteFindMax: empty map has no maximal element") . maxViewWithKey
1101
1102 -- | /O(min(n,W))/. Delete and find the minimal element.
1103 deleteFindMin :: IntMap a -> ((Key, a), IntMap a)
1104 deleteFindMin = fromMaybe (error "deleteFindMin: empty map has no minimal element") . minViewWithKey
1105
1106 -- | /O(min(n,W))/. The minimal key of the map.
1107 findMin :: IntMap a -> (Key, a)
1108 findMin Nil = error $ "findMin: empty map has no minimal element"
1109 findMin (Tip k v) = (k,v)
1110 findMin (Bin _ m l r)
1111 | m < 0 = go r
1112 | otherwise = go l
1113 where go (Tip k v) = (k,v)
1114 go (Bin _ _ l' _) = go l'
1115 go Nil = error "findMax Nil"
1116
1117 -- | /O(min(n,W))/. The maximal key of the map.
1118 findMax :: IntMap a -> (Key, a)
1119 findMax Nil = error $ "findMax: empty map has no maximal element"
1120 findMax (Tip k v) = (k,v)
1121 findMax (Bin _ m l r)
1122 | m < 0 = go l
1123 | otherwise = go r
1124 where go (Tip k v) = (k,v)
1125 go (Bin _ _ _ r') = go r'
1126 go Nil = error "findMax Nil"
1127
1128 -- | /O(min(n,W))/. Delete the minimal key. Returns an empty map if the map is empty.
1129 --
1130 -- Note that this is a change of behaviour for consistency with 'Data.Map.Map' &#8211;
1131 -- versions prior to 0.5 threw an error if the 'IntMap' was already empty.
1132 deleteMin :: IntMap a -> IntMap a
1133 deleteMin = maybe Nil snd . minView
1134
1135 -- | /O(min(n,W))/. Delete the maximal key. Returns an empty map if the map is empty.
1136 --
1137 -- Note that this is a change of behaviour for consistency with 'Data.Map.Map' &#8211;
1138 -- versions prior to 0.5 threw an error if the 'IntMap' was already empty.
1139 deleteMax :: IntMap a -> IntMap a
1140 deleteMax = maybe Nil snd . maxView
1141
1142
1143 {--------------------------------------------------------------------
1144 Submap
1145 --------------------------------------------------------------------}
1146 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
1147 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
1148 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
1149 isProperSubmapOf m1 m2
1150 = isProperSubmapOfBy (==) m1 m2
1151
1152 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
1153 The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
1154 @m1@ and @m2@ are not equal,
1155 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
1156 applied to their respective values. For example, the following
1157 expressions are all 'True':
1158
1159 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1160 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1161
1162 But the following are all 'False':
1163
1164 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
1165 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
1166 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1167 -}
1168 isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
1169 isProperSubmapOfBy predicate t1 t2
1170 = case submapCmp predicate t1 t2 of
1171 LT -> True
1172 _ -> False
1173
1174 submapCmp :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Ordering
1175 submapCmp predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1176 | shorter m1 m2 = GT
1177 | shorter m2 m1 = submapCmpLt
1178 | p1 == p2 = submapCmpEq
1179 | otherwise = GT -- disjoint
1180 where
1181 submapCmpLt | nomatch p1 p2 m2 = GT
1182 | zero p1 m2 = submapCmp predicate t1 l2
1183 | otherwise = submapCmp predicate t1 r2
1184 submapCmpEq = case (submapCmp predicate l1 l2, submapCmp predicate r1 r2) of
1185 (GT,_ ) -> GT
1186 (_ ,GT) -> GT
1187 (EQ,EQ) -> EQ
1188 _ -> LT
1189
1190 submapCmp _ (Bin _ _ _ _) _ = GT
1191 submapCmp predicate (Tip kx x) (Tip ky y)
1192 | (kx == ky) && predicate x y = EQ
1193 | otherwise = GT -- disjoint
1194 submapCmp predicate (Tip k x) t
1195 = case lookup k t of
1196 Just y | predicate x y -> LT
1197 _ -> GT -- disjoint
1198 submapCmp _ Nil Nil = EQ
1199 submapCmp _ Nil _ = LT
1200
1201 -- | /O(n+m)/. Is this a submap?
1202 -- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
1203 isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
1204 isSubmapOf m1 m2
1205 = isSubmapOfBy (==) m1 m2
1206
1207 {- | /O(n+m)/.
1208 The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if
1209 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
1210 applied to their respective values. For example, the following
1211 expressions are all 'True':
1212
1213 > isSubmapOfBy (==) (fromList [(1,1)]) (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),(2,2)])
1216
1217 But the following are all 'False':
1218
1219 > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)])
1220 > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1221 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
1222 -}
1223 isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
1224 isSubmapOfBy predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1225 | shorter m1 m2 = False
1226 | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy predicate t1 l2
1227 else isSubmapOfBy predicate t1 r2)
1228 | otherwise = (p1==p2) && isSubmapOfBy predicate l1 l2 && isSubmapOfBy predicate r1 r2
1229 isSubmapOfBy _ (Bin _ _ _ _) _ = False
1230 isSubmapOfBy predicate (Tip k x) t = case lookup k t of
1231 Just y -> predicate x y
1232 Nothing -> False
1233 isSubmapOfBy _ Nil _ = True
1234
1235 {--------------------------------------------------------------------
1236 Mapping
1237 --------------------------------------------------------------------}
1238 -- | /O(n)/. Map a function over all values in the map.
1239 --
1240 -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")]
1241
1242 map :: (a -> b) -> IntMap a -> IntMap b
1243 map f t
1244 = case t of
1245 Bin p m l r -> Bin p m (map f l) (map f r)
1246 Tip k x -> Tip k (f x)
1247 Nil -> Nil
1248
1249 -- | /O(n)/. Map a function over all values in the map.
1250 --
1251 -- > let f key x = (show key) ++ ":" ++ x
1252 -- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")]
1253
1254 mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b
1255 mapWithKey f t
1256 = case t of
1257 Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r)
1258 Tip k x -> Tip k (f k x)
1259 Nil -> Nil
1260
1261 -- | /O(n)/.
1262 -- @'traverseWithKey' f s == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@
1263 -- That is, behaves exactly like a regular 'traverse' except that the traversing
1264 -- function also has access to the key associated with a value.
1265 --
1266 -- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(1, 'a'), (5, 'e')]) == Just (fromList [(1, 'b'), (5, 'f')])
1267 -- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(2, 'c')]) == Nothing
1268 {-# INLINE traverseWithKey #-}
1269 traverseWithKey :: Applicative t => (Key -> a -> t b) -> IntMap a -> t (IntMap b)
1270 traverseWithKey f = go
1271 where
1272 go Nil = pure Nil
1273 go (Tip k v) = Tip k <$> f k v
1274 go (Bin p m l r) = Bin p m <$> go l <*> go r
1275
1276 -- | /O(n)/. The function @'mapAccum'@ threads an accumulating
1277 -- argument through the map in ascending order of keys.
1278 --
1279 -- > let f a b = (a ++ b, b ++ "X")
1280 -- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")])
1281
1282 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1283 mapAccum f = mapAccumWithKey (\a' _ x -> f a' x)
1284
1285 -- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating
1286 -- argument through the map in ascending order of keys.
1287 --
1288 -- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X")
1289 -- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")])
1290
1291 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1292 mapAccumWithKey f a t
1293 = mapAccumL f a t
1294
1295 -- | /O(n)/. The function @'mapAccumL'@ threads an accumulating
1296 -- argument through the map in ascending order of keys.
1297 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1298 mapAccumL f a t
1299 = case t of
1300 Bin p m l r -> let (a1,l') = mapAccumL f a l
1301 (a2,r') = mapAccumL f a1 r
1302 in (a2,Bin p m l' r')
1303 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
1304 Nil -> (a,Nil)
1305
1306 -- | /O(n)/. The function @'mapAccumR'@ threads an accumulating
1307 -- argument through the map in descending order of keys.
1308 mapAccumRWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1309 mapAccumRWithKey f a t
1310 = case t of
1311 Bin p m l r -> let (a1,r') = mapAccumRWithKey f a r
1312 (a2,l') = mapAccumRWithKey f a1 l
1313 in (a2,Bin p m l' r')
1314 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
1315 Nil -> (a,Nil)
1316
1317 -- | /O(n*min(n,W))/.
1318 -- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@.
1319 --
1320 -- The size of the result may be smaller if @f@ maps two or more distinct
1321 -- keys to the same new key. In this case the value at the greatest of the
1322 -- original keys is retained.
1323 --
1324 -- > mapKeys (+ 1) (fromList [(5,"a"), (3,"b")]) == fromList [(4, "b"), (6, "a")]
1325 -- > mapKeys (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "c"
1326 -- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c"
1327
1328 mapKeys :: (Key->Key) -> IntMap a -> IntMap a
1329 mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) []
1330
1331 -- | /O(n*min(n,W))/.
1332 -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@.
1333 --
1334 -- The size of the result may be smaller if @f@ maps two or more distinct
1335 -- keys to the same new key. In this case the associated values will be
1336 -- combined using @c@.
1337 --
1338 -- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab"
1339 -- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab"
1340
1341 mapKeysWith :: (a -> a -> a) -> (Key->Key) -> IntMap a -> IntMap a
1342 mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) []
1343
1344 -- | /O(n*min(n,W))/.
1345 -- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@
1346 -- is strictly monotonic.
1347 -- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@.
1348 -- /The precondition is not checked./
1349 -- Semi-formally, we have:
1350 --
1351 -- > and [x < y ==> f x < f y | x <- ls, y <- ls]
1352 -- > ==> mapKeysMonotonic f s == mapKeys f s
1353 -- > where ls = keys s
1354 --
1355 -- This means that @f@ maps distinct original keys to distinct resulting keys.
1356 -- This function has slightly better performance than 'mapKeys'.
1357 --
1358 -- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")]
1359
1360 mapKeysMonotonic :: (Key->Key) -> IntMap a -> IntMap a
1361 mapKeysMonotonic f = fromDistinctAscList . foldrWithKey (\k x xs -> (f k, x) : xs) []
1362
1363 {--------------------------------------------------------------------
1364 Filter
1365 --------------------------------------------------------------------}
1366 -- | /O(n)/. Filter all values that satisfy some predicate.
1367 --
1368 -- > filter (> "a") (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1369 -- > filter (> "x") (fromList [(5,"a"), (3,"b")]) == empty
1370 -- > filter (< "a") (fromList [(5,"a"), (3,"b")]) == empty
1371
1372 filter :: (a -> Bool) -> IntMap a -> IntMap a
1373 filter p m
1374 = filterWithKey (\_ x -> p x) m
1375
1376 -- | /O(n)/. Filter all keys\/values that satisfy some predicate.
1377 --
1378 -- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1379
1380 filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a
1381 filterWithKey predicate t
1382 = case t of
1383 Bin p m l r
1384 -> bin p m (filterWithKey predicate l) (filterWithKey predicate r)
1385 Tip k x
1386 | predicate k x -> t
1387 | otherwise -> Nil
1388 Nil -> Nil
1389
1390 -- | /O(n)/. Partition the map according to some predicate. The first
1391 -- map contains all elements that satisfy the predicate, the second all
1392 -- elements that fail the predicate. See also 'split'.
1393 --
1394 -- > partition (> "a") (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
1395 -- > partition (< "x") (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
1396 -- > partition (> "x") (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
1397
1398 partition :: (a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
1399 partition p m
1400 = partitionWithKey (\_ x -> p x) m
1401
1402 -- | /O(n)/. Partition the map according to some predicate. The first
1403 -- map contains all elements that satisfy the predicate, the second all
1404 -- elements that fail the predicate. See also 'split'.
1405 --
1406 -- > partitionWithKey (\ k _ -> k > 3) (fromList [(5,"a"), (3,"b")]) == (singleton 5 "a", singleton 3 "b")
1407 -- > partitionWithKey (\ k _ -> k < 7) (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
1408 -- > partitionWithKey (\ k _ -> k > 7) (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
1409
1410 partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
1411 partitionWithKey predicate0 t0 = toPair $ go predicate0 t0
1412 where
1413 go predicate t
1414 = case t of
1415 Bin p m l r
1416 -> let (l1 :*: l2) = go predicate l
1417 (r1 :*: r2) = go predicate r
1418 in bin p m l1 r1 :*: bin p m l2 r2
1419 Tip k x
1420 | predicate k x -> (t :*: Nil)
1421 | otherwise -> (Nil :*: t)
1422 Nil -> (Nil :*: Nil)
1423
1424 -- | /O(n)/. Map values and collect the 'Just' results.
1425 --
1426 -- > let f x = if x == "a" then Just "new a" else Nothing
1427 -- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a"
1428
1429 mapMaybe :: (a -> Maybe b) -> IntMap a -> IntMap b
1430 mapMaybe f = mapMaybeWithKey (\_ x -> f x)
1431
1432 -- | /O(n)/. Map keys\/values and collect the 'Just' results.
1433 --
1434 -- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing
1435 -- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3"
1436
1437 mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b
1438 mapMaybeWithKey f (Bin p m l r)
1439 = bin p m (mapMaybeWithKey f l) (mapMaybeWithKey f r)
1440 mapMaybeWithKey f (Tip k x) = case f k x of
1441 Just y -> Tip k y
1442 Nothing -> Nil
1443 mapMaybeWithKey _ Nil = Nil
1444
1445 -- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
1446 --
1447 -- > let f a = if a < "c" then Left a else Right a
1448 -- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1449 -- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")])
1450 -- >
1451 -- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1452 -- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1453
1454 mapEither :: (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
1455 mapEither f m
1456 = mapEitherWithKey (\_ x -> f x) m
1457
1458 -- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
1459 --
1460 -- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a)
1461 -- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1462 -- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")])
1463 -- >
1464 -- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1465 -- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")])
1466
1467 mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
1468 mapEitherWithKey f0 t0 = toPair $ go f0 t0
1469 where
1470 go f (Bin p m l r)
1471 = bin p m l1 r1 :*: bin p m l2 r2
1472 where
1473 (l1 :*: l2) = go f l
1474 (r1 :*: r2) = go f r
1475 go f (Tip k x) = case f k x of
1476 Left y -> (Tip k y :*: Nil)
1477 Right z -> (Nil :*: Tip k z)
1478 go _ Nil = (Nil :*: Nil)
1479
1480 -- | /O(min(n,W))/. The expression (@'split' k map@) is a pair @(map1,map2)@
1481 -- where all keys in @map1@ are lower than @k@ and all keys in
1482 -- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
1483 --
1484 -- > split 2 (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3,"b"), (5,"a")])
1485 -- > split 3 (fromList [(5,"a"), (3,"b")]) == (empty, singleton 5 "a")
1486 -- > split 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
1487 -- > split 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", empty)
1488 -- > split 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], empty)
1489
1490 split :: Key -> IntMap a -> (IntMap a, IntMap a)
1491 split k t =
1492 case t of
1493 Bin _ m l r
1494 | m < 0 -> if k >= 0 -- handle negative numbers.
1495 then case go k l of (lt :*: gt) -> let lt' = union r lt
1496 in lt' `seq` (lt', gt)
1497 else case go k r of (lt :*: gt) -> let gt' = union gt l
1498 in gt' `seq` (lt, gt')
1499 _ -> case go k t of
1500 (lt :*: gt) -> (lt, gt)
1501 where
1502 go k' t'@(Bin p m l r) | nomatch k' p m = if k' > p then t' :*: Nil else Nil :*: t'
1503 | zero k' m = case go k' l of (lt :*: gt) -> lt :*: union gt r
1504 | otherwise = case go k' r of (lt :*: gt) -> union l lt :*: gt
1505 go k' t'@(Tip ky _) | k' > ky = (t' :*: Nil)
1506 | k' < ky = (Nil :*: t')
1507 | otherwise = (Nil :*: Nil)
1508 go _ Nil = (Nil :*: Nil)
1509
1510 -- | /O(min(n,W))/. Performs a 'split' but also returns whether the pivot
1511 -- key was found in the original map.
1512 --
1513 -- > splitLookup 2 (fromList [(5,"a"), (3,"b")]) == (empty, Nothing, fromList [(3,"b"), (5,"a")])
1514 -- > splitLookup 3 (fromList [(5,"a"), (3,"b")]) == (empty, Just "b", singleton 5 "a")
1515 -- > splitLookup 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Nothing, singleton 5 "a")
1516 -- > splitLookup 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Just "a", empty)
1517 -- > splitLookup 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], Nothing, empty)
1518
1519 splitLookup :: Key -> IntMap a -> (IntMap a, Maybe a, IntMap a)
1520 splitLookup k t =
1521 case t of
1522 Bin _ m l r
1523 | m < 0 -> if k >= 0 -- handle negative numbers.
1524 then case go k l of
1525 (lt, fnd, gt) -> let lt' = union r lt
1526 in lt' `seq` (lt', fnd, gt)
1527 else case go k r of
1528 (lt, fnd, gt) -> let gt' = union gt l
1529 in gt' `seq` (lt, fnd, gt')
1530 _ -> go k t
1531 where
1532 go k' t'@(Bin p m l r)
1533 | nomatch k' p m = if k' > p then (t', Nothing, Nil) else (Nil, Nothing, t')
1534 | zero k' m = case go k' l of
1535 (lt, fnd, gt) -> let gt' = union gt r in gt' `seq` (lt, fnd, gt')
1536 | otherwise = case go k' r of
1537 (lt, fnd, gt) -> let lt' = union l lt in lt' `seq` (lt', fnd, gt)
1538 go k' t'@(Tip ky y) | k' > ky = (t', Nothing, Nil)
1539 | k' < ky = (Nil, Nothing, t')
1540 | otherwise = (Nil, Just y, Nil)
1541 go _ Nil = (Nil, Nothing, Nil)
1542
1543 {--------------------------------------------------------------------
1544 Fold
1545 --------------------------------------------------------------------}
1546 -- | /O(n)/. Fold the values in the map using the given right-associative
1547 -- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'elems'@.
1548 --
1549 -- For example,
1550 --
1551 -- > elems map = foldr (:) [] map
1552 --
1553 -- > let f a len = len + (length a)
1554 -- > foldr f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
1555 foldr :: (a -> b -> b) -> b -> IntMap a -> b
1556 foldr f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1557 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1558 | otherwise -> go (go z r) l
1559 _ -> go z t
1560 where
1561 go z' Nil = z'
1562 go z' (Tip _ x) = f x z'
1563 go z' (Bin _ _ l r) = go (go z' r) l
1564 {-# INLINE foldr #-}
1565
1566 -- | /O(n)/. A strict version of 'foldr'. Each application of the operator is
1567 -- evaluated before using the result in the next application. This
1568 -- function is strict in the starting value.
1569 foldr' :: (a -> b -> b) -> b -> IntMap a -> b
1570 foldr' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1571 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1572 | otherwise -> go (go z r) l
1573 _ -> go z t
1574 where
1575 STRICT_1_OF_2(go)
1576 go z' Nil = z'
1577 go z' (Tip _ x) = f x z'
1578 go z' (Bin _ _ l r) = go (go z' r) l
1579 {-# INLINE foldr' #-}
1580
1581 -- | /O(n)/. Fold the values in the map using the given left-associative
1582 -- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@.
1583 --
1584 -- For example,
1585 --
1586 -- > elems = reverse . foldl (flip (:)) []
1587 --
1588 -- > let f len a = len + (length a)
1589 -- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
1590 foldl :: (a -> b -> a) -> a -> IntMap b -> a
1591 foldl f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1592 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1593 | otherwise -> go (go z l) r
1594 _ -> go z t
1595 where
1596 go z' Nil = z'
1597 go z' (Tip _ x) = f z' x
1598 go z' (Bin _ _ l r) = go (go z' l) r
1599 {-# INLINE foldl #-}
1600
1601 -- | /O(n)/. A strict version of 'foldl'. Each application of the operator is
1602 -- evaluated before using the result in the next application. This
1603 -- function is strict in the starting value.
1604 foldl' :: (a -> b -> a) -> a -> IntMap b -> a
1605 foldl' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1606 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1607 | otherwise -> go (go z l) r
1608 _ -> go z t
1609 where
1610 STRICT_1_OF_2(go)
1611 go z' Nil = z'
1612 go z' (Tip _ x) = f z' x
1613 go z' (Bin _ _ l r) = go (go z' l) r
1614 {-# INLINE foldl' #-}
1615
1616 -- | /O(n)/. Fold the keys and values in the map using the given right-associative
1617 -- binary operator, such that
1618 -- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
1619 --
1620 -- For example,
1621 --
1622 -- > keys map = foldrWithKey (\k x ks -> k:ks) [] map
1623 --
1624 -- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
1625 -- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)"
1626 foldrWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
1627 foldrWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1628 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1629 | otherwise -> go (go z r) l
1630 _ -> go z t
1631 where
1632 go z' Nil = z'
1633 go z' (Tip kx x) = f kx x z'
1634 go z' (Bin _ _ l r) = go (go z' r) l
1635 {-# INLINE foldrWithKey #-}
1636
1637 -- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is
1638 -- evaluated before using the result in the next application. This
1639 -- function is strict in the starting value.
1640 foldrWithKey' :: (Key -> a -> b -> b) -> b -> IntMap a -> b
1641 foldrWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1642 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1643 | otherwise -> go (go z r) l
1644 _ -> go z t
1645 where
1646 STRICT_1_OF_2(go)
1647 go z' Nil = z'
1648 go z' (Tip kx x) = f kx x z'
1649 go z' (Bin _ _ l r) = go (go z' r) l
1650 {-# INLINE foldrWithKey' #-}
1651
1652 -- | /O(n)/. Fold the keys and values in the map using the given left-associative
1653 -- binary operator, such that
1654 -- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@.
1655 --
1656 -- For example,
1657 --
1658 -- > keys = reverse . foldlWithKey (\ks k x -> k:ks) []
1659 --
1660 -- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
1661 -- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)"
1662 foldlWithKey :: (a -> Key -> b -> a) -> a -> IntMap b -> a
1663 foldlWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1664 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1665 | otherwise -> go (go z l) r
1666 _ -> go z t
1667 where
1668 go z' Nil = z'
1669 go z' (Tip kx x) = f z' kx x
1670 go z' (Bin _ _ l r) = go (go z' l) r
1671 {-# INLINE foldlWithKey #-}
1672
1673 -- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is
1674 -- evaluated before using the result in the next application. This
1675 -- function is strict in the starting value.
1676 foldlWithKey' :: (a -> Key -> b -> a) -> a -> IntMap b -> a
1677 foldlWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1678 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1679 | otherwise -> go (go z l) r
1680 _ -> go z t
1681 where
1682 STRICT_1_OF_2(go)
1683 go z' Nil = z'
1684 go z' (Tip kx x) = f z' kx x
1685 go z' (Bin _ _ l r) = go (go z' l) r
1686 {-# INLINE foldlWithKey' #-}
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 Finding the highest bit (mask) in a word [x] can be done efficiently in
2064 three ways:
2065 * convert to a floating point value and the mantissa tells us the
2066 [log2(x)] that corresponds with the highest bit position. The mantissa
2067 is retrieved either via the standard C function [frexp] or by some bit
2068 twiddling on IEEE compatible numbers (float). Note that one needs to
2069 use at least [double] precision for an accurate mantissa of 32 bit
2070 numbers.
2071 * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
2072 * use processor specific assembler instruction (asm).
2073
2074 The most portable way would be [bit], but is it efficient enough?
2075 I have measured the cycle counts of the different methods on an AMD
2076 Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
2077
2078 highestBitMask: method cycles
2079 --------------
2080 frexp 200
2081 float 33
2082 bit 11
2083 asm 12
2084
2085 highestBit: method cycles
2086 --------------
2087 frexp 195
2088 float 33
2089 bit 11
2090 asm 11
2091
2092 Wow, the bit twiddling is on today's RISC like machines even faster
2093 than a single CISC instruction (BSR)!
2094 ----------------------------------------------------------------------}
2095
2096 {----------------------------------------------------------------------
2097 [highestBitMask] returns a word where only the highest bit is set.
2098 It is found by first setting all bits in lower positions than the
2099 highest bit and than taking an exclusive or with the original value.
2100 Allthough the function may look expensive, GHC compiles this into
2101 excellent C code that subsequently compiled into highly efficient
2102 machine code. The algorithm is derived from Jorg Arndt's FXT library.
2103 ----------------------------------------------------------------------}
2104 highestBitMask :: Nat -> Nat
2105 highestBitMask x0
2106 = case (x0 .|. shiftRL x0 1) of
2107 x1 -> case (x1 .|. shiftRL x1 2) of
2108 x2 -> case (x2 .|. shiftRL x2 4) of
2109 x3 -> case (x3 .|. shiftRL x3 8) of
2110 x4 -> case (x4 .|. shiftRL x4 16) of
2111 #if !(defined(__GLASGOW_HASKELL__) && WORD_SIZE_IN_BITS==32)
2112 x5 -> case (x5 .|. shiftRL x5 32) of -- for 64 bit platforms
2113 #endif
2114 x6 -> (x6 `xor` (shiftRL x6 1))
2115 {-# INLINE highestBitMask #-}
2116
2117
2118 {--------------------------------------------------------------------
2119 Utilities
2120 --------------------------------------------------------------------}
2121
2122 foldlStrict :: (a -> b -> a) -> a -> [b] -> a
2123 foldlStrict f = go
2124 where
2125 go z [] = z
2126 go z (x:xs) = let z' = f z x in z' `seq` go z' xs
2127 {-# INLINE foldlStrict #-}
2128
2129 {--------------------------------------------------------------------
2130 Debugging
2131 --------------------------------------------------------------------}
2132 -- | /O(n)/. Show the tree that implements the map. The tree is shown
2133 -- in a compressed, hanging format.
2134 showTree :: Show a => IntMap a -> String
2135 showTree s
2136 = showTreeWith True False s
2137
2138
2139 {- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows
2140 the tree that implements the map. If @hang@ is
2141 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
2142 @wide@ is 'True', an extra wide version is shown.
2143 -}
2144 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
2145 showTreeWith hang wide t
2146 | hang = (showsTreeHang wide [] t) ""
2147 | otherwise = (showsTree wide [] [] t) ""
2148
2149 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
2150 showsTree wide lbars rbars t
2151 = case t of
2152 Bin p m l r
2153 -> showsTree wide (withBar rbars) (withEmpty rbars) r .
2154 showWide wide rbars .
2155 showsBars lbars . showString (showBin p m) . showString "\n" .
2156 showWide wide lbars .
2157 showsTree wide (withEmpty lbars) (withBar lbars) l
2158 Tip k x
2159 -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n"
2160 Nil -> showsBars lbars . showString "|\n"
2161
2162 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
2163 showsTreeHang wide bars t
2164 = case t of
2165 Bin p m l r
2166 -> showsBars bars . showString (showBin p m) . showString "\n" .
2167 showWide wide bars .
2168 showsTreeHang wide (withBar bars) l .
2169 showWide wide bars .
2170 showsTreeHang wide (withEmpty bars) r
2171 Tip k x
2172 -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n"
2173 Nil -> showsBars bars . showString "|\n"
2174
2175 showBin :: Prefix -> Mask -> String
2176 showBin _ _
2177 = "*" -- ++ show (p,m)
2178
2179 showWide :: Bool -> [String] -> String -> String
2180 showWide wide bars
2181 | wide = showString (concat (reverse bars)) . showString "|\n"
2182 | otherwise = id
2183
2184 showsBars :: [String] -> ShowS
2185 showsBars bars
2186 = case bars of
2187 [] -> id
2188 _ -> showString (concat (reverse (tail bars))) . showString node
2189
2190 node :: String
2191 node = "+--"
2192
2193 withBar, withEmpty :: [String] -> [String]
2194 withBar bars = "| ":bars
2195 withEmpty bars = " ":bars