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