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