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