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