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