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