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