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