Make strict IntMap merges strict
[packages/containers.git] / Data / IntMap / Internal.hs
1 {-# LANGUAGE CPP #-}
2 {-# LANGUAGE BangPatterns #-}
3 {-# LANGUAGE PatternGuards #-}
4 #if __GLASGOW_HASKELL__
5 {-# LANGUAGE MagicHash, DeriveDataTypeable, StandaloneDeriving #-}
6 {-# LANGUAGE ScopedTypeVariables #-}
7 #endif
8 #if !defined(TESTING) && defined(__GLASGOW_HASKELL__)
9 {-# LANGUAGE Trustworthy #-}
10 #endif
11 #if __GLASGOW_HASKELL__ >= 708
12 {-# LANGUAGE TypeFamilies #-}
13 #endif
14
15 {-# OPTIONS_HADDOCK not-home #-}
16
17 #include "containers.h"
18
19 -----------------------------------------------------------------------------
20 -- |
21 -- Module : Data.IntMap.Internal
22 -- Copyright : (c) Daan Leijen 2002
23 -- (c) Andriy Palamarchuk 2008
24 -- (c) wren romano 2016
25 -- License : BSD-style
26 -- Maintainer : libraries@haskell.org
27 -- Portability : portable
28 --
29 -- = WARNING
30 --
31 -- This module is considered __internal__.
32 --
33 -- The Package Versioning Policy __does not apply__.
34 --
35 -- The contents of this module may change __in any way whatsoever__
36 -- and __without any warning__ between minor versions of this package.
37 --
38 -- Authors importing this module are expected to track development
39 -- closely.
40 --
41 -- = Description
42 --
43 -- This defines the data structures and core (hidden) manipulations
44 -- on representations.
45 --
46 -- @since 0.5.9
47 -----------------------------------------------------------------------------
48
49 -- [Note: INLINE bit fiddling]
50 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~
51 -- It is essential that the bit fiddling functions like mask, zero, branchMask
52 -- etc are inlined. If they do not, the memory allocation skyrockets. The GHC
53 -- usually gets it right, but it is disastrous if it does not. Therefore we
54 -- explicitly mark these functions INLINE.
55
56
57 -- [Note: Local 'go' functions and capturing]
58 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
59 -- Care must be taken when using 'go' function which captures an argument.
60 -- Sometimes (for example when the argument is passed to a data constructor,
61 -- as in insert), GHC heap-allocates more than necessary. Therefore C-- code
62 -- must be checked for increased allocation when creating and modifying such
63 -- functions.
64
65
66 -- [Note: Order of constructors]
67 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
68 -- The order of constructors of IntMap matters when considering performance.
69 -- Currently in GHC 7.0, when type has 3 constructors, they are matched from
70 -- the first to the last -- the best performance is achieved when the
71 -- constructors are ordered by frequency.
72 -- On GHC 7.0, reordering constructors from Nil | Tip | Bin to Bin | Tip | Nil
73 -- improves the benchmark by circa 10%.
74
75 module Data.IntMap.Internal (
76 -- * Map type
77 IntMap(..), Key -- instance Eq,Show
78
79 -- * Operators
80 , (!), (!?), (\\)
81
82 -- * Query
83 , null
84 , size
85 , member
86 , notMember
87 , lookup
88 , findWithDefault
89 , lookupLT
90 , lookupGT
91 , lookupLE
92 , lookupGE
93
94 -- * Construction
95 , empty
96 , singleton
97
98 -- ** Insertion
99 , insert
100 , insertWith
101 , insertWithKey
102 , insertLookupWithKey
103
104 -- ** Delete\/Update
105 , delete
106 , adjust
107 , adjustWithKey
108 , update
109 , updateWithKey
110 , updateLookupWithKey
111 , alter
112 , alterF
113
114 -- * Combine
115
116 -- ** Union
117 , union
118 , unionWith
119 , unionWithKey
120 , unions
121 , unionsWith
122
123 -- ** Difference
124 , difference
125 , differenceWith
126 , differenceWithKey
127
128 -- ** Intersection
129 , intersection
130 , intersectionWith
131 , intersectionWithKey
132
133 -- ** General combining function
134 , SimpleWhenMissing
135 , SimpleWhenMatched
136 , runWhenMatched
137 , runWhenMissing
138 , merge
139 -- *** @WhenMatched@ tactics
140 , zipWithMaybeMatched
141 , zipWithMatched
142 -- *** @WhenMissing@ tactics
143 , mapMaybeMissing
144 , dropMissing
145 , preserveMissing
146 , mapMissing
147 , filterMissing
148
149 -- ** Applicative general combining function
150 , WhenMissing (..)
151 , WhenMatched (..)
152 , mergeA
153 -- *** @WhenMatched@ tactics
154 -- | The tactics described for 'merge' work for
155 -- 'mergeA' as well. Furthermore, the following
156 -- are available.
157 , zipWithMaybeAMatched
158 , zipWithAMatched
159 -- *** @WhenMissing@ tactics
160 -- | The tactics described for 'merge' work for
161 -- 'mergeA' as well. Furthermore, the following
162 -- are available.
163 , traverseMaybeMissing
164 , traverseMissing
165 , filterAMissing
166
167 -- ** Deprecated general combining function
168 , mergeWithKey
169 , mergeWithKey'
170
171 -- * Traversal
172 -- ** Map
173 , map
174 , mapWithKey
175 , traverseWithKey
176 , traverseMaybeWithKey
177 , mapAccum
178 , mapAccumWithKey
179 , mapAccumRWithKey
180 , mapKeys
181 , mapKeysWith
182 , mapKeysMonotonic
183
184 -- * Folds
185 , foldr
186 , foldl
187 , foldrWithKey
188 , foldlWithKey
189 , foldMapWithKey
190
191 -- ** Strict folds
192 , foldr'
193 , foldl'
194 , foldrWithKey'
195 , foldlWithKey'
196
197 -- * Conversion
198 , elems
199 , keys
200 , assocs
201 , keysSet
202 , fromSet
203
204 -- ** Lists
205 , toList
206 , fromList
207 , fromListWith
208 , fromListWithKey
209
210 -- ** Ordered lists
211 , toAscList
212 , toDescList
213 , fromAscList
214 , fromAscListWith
215 , fromAscListWithKey
216 , fromDistinctAscList
217
218 -- * Filter
219 , filter
220 , filterWithKey
221 , restrictKeys
222 , withoutKeys
223 , partition
224 , partitionWithKey
225
226 , mapMaybe
227 , mapMaybeWithKey
228 , mapEither
229 , mapEitherWithKey
230
231 , split
232 , splitLookup
233 , splitRoot
234
235 -- * Submap
236 , isSubmapOf, isSubmapOfBy
237 , isProperSubmapOf, isProperSubmapOfBy
238
239 -- * Min\/Max
240 , lookupMin
241 , lookupMax
242 , findMin
243 , findMax
244 , deleteMin
245 , deleteMax
246 , deleteFindMin
247 , deleteFindMax
248 , updateMin
249 , updateMax
250 , updateMinWithKey
251 , updateMaxWithKey
252 , minView
253 , maxView
254 , minViewWithKey
255 , maxViewWithKey
256
257 -- * Debugging
258 , showTree
259 , showTreeWith
260
261 -- * Internal types
262 , Mask, Prefix, Nat
263
264 -- * Utility
265 , natFromInt
266 , intFromNat
267 , link
268 , bin
269 , binCheckLeft
270 , binCheckRight
271 , zero
272 , nomatch
273 , match
274 , mask
275 , maskW
276 , shorter
277 , branchMask
278 , highestBitMask
279
280 -- * Used by "IntMap.Merge.Lazy" and "IntMap.Merge.Strict"
281 , mapWhenMissing
282 , mapWhenMatched
283 , lmapWhenMissing
284 , contramapFirstWhenMatched
285 , contramapSecondWhenMatched
286 , mapGentlyWhenMissing
287 , mapGentlyWhenMatched
288 ) where
289
290 #if MIN_VERSION_base(4,8,0)
291 import Data.Functor.Identity (Identity (..))
292 import Control.Applicative (liftA2)
293 #else
294 import Control.Applicative (Applicative(pure, (<*>)), (<$>), liftA2)
295 import Data.Monoid (Monoid(..))
296 import Data.Traversable (Traversable(traverse))
297 import Data.Word (Word)
298 #endif
299 #if MIN_VERSION_base(4,9,0)
300 import Data.Semigroup (Semigroup(stimes))
301 #endif
302 #if !(MIN_VERSION_base(4,11,0)) && MIN_VERSION_base(4,9,0)
303 import Data.Semigroup (Semigroup((<>)))
304 #endif
305 #if MIN_VERSION_base(4,9,0)
306 import Data.Semigroup (stimesIdempotentMonoid)
307 import Data.Functor.Classes
308 #endif
309
310 import Control.DeepSeq (NFData(rnf))
311 import Data.Bits
312 import qualified Data.Foldable as Foldable
313 #if !MIN_VERSION_base(4,8,0)
314 import Data.Foldable (Foldable())
315 #endif
316 import Data.Maybe (fromMaybe)
317 import Data.Typeable
318 import Prelude hiding (lookup, map, filter, foldr, foldl, null)
319
320 import Data.IntSet.Internal (Key)
321 import qualified Data.IntSet.Internal as IntSet
322 import Utils.Containers.Internal.BitUtil
323 import Utils.Containers.Internal.StrictPair
324
325 #if __GLASGOW_HASKELL__
326 import Data.Data (Data(..), Constr, mkConstr, constrIndex, Fixity(Prefix),
327 DataType, mkDataType)
328 import GHC.Exts (build)
329 #if !MIN_VERSION_base(4,8,0)
330 import Data.Functor ((<$))
331 #endif
332 #if __GLASGOW_HASKELL__ >= 708
333 import qualified GHC.Exts as GHCExts
334 #endif
335 import Text.Read
336 #endif
337 import qualified Control.Category as Category
338 #if __GLASGOW_HASKELL__ >= 709
339 import Data.Coerce
340 #endif
341
342
343 -- A "Nat" is a natural machine word (an unsigned Int)
344 type Nat = Word
345
346 natFromInt :: Key -> Nat
347 natFromInt = fromIntegral
348 {-# INLINE natFromInt #-}
349
350 intFromNat :: Nat -> Key
351 intFromNat = fromIntegral
352 {-# INLINE intFromNat #-}
353
354 {--------------------------------------------------------------------
355 Types
356 --------------------------------------------------------------------}
357
358
359 -- | A map of integers to values @a@.
360
361 -- See Note: Order of constructors
362 data IntMap a = Bin {-# UNPACK #-} !Prefix
363 {-# UNPACK #-} !Mask
364 !(IntMap a)
365 !(IntMap a)
366 -- Fields:
367 -- prefix: The most significant bits shared by all keys in this Bin.
368 -- mask: The switching bit to determine if a key should follow the left
369 -- or right subtree of a 'Bin'.
370 -- Invariant: Nil is never found as a child of Bin.
371 -- Invariant: The Mask is a power of 2. It is the largest bit position at which
372 -- two keys of the map differ.
373 -- Invariant: Prefix is the common high-order bits that all elements share to
374 -- the left of the Mask bit.
375 -- Invariant: In Bin prefix mask left right, left consists of the elements that
376 -- don't have the mask bit set; right is all the elements that do.
377 | Tip {-# UNPACK #-} !Key a
378 | Nil
379
380 type Prefix = Int
381 type Mask = Int
382
383
384 -- Some stuff from "Data.IntSet.Internal", for 'restrictKeys' and
385 -- 'withoutKeys' to use.
386 type IntSetPrefix = Int
387 type IntSetBitMap = Word
388
389 bitmapOf :: Int -> IntSetBitMap
390 bitmapOf x = shiftLL 1 (x .&. IntSet.suffixBitMask)
391 {-# INLINE bitmapOf #-}
392
393 {--------------------------------------------------------------------
394 Operators
395 --------------------------------------------------------------------}
396
397 -- | /O(min(n,W))/. Find the value at a key.
398 -- Calls 'error' when the element can not be found.
399 --
400 -- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map
401 -- > fromList [(5,'a'), (3,'b')] ! 5 == 'a'
402
403 (!) :: IntMap a -> Key -> a
404 (!) m k = find k m
405
406 -- | /O(min(n,W))/. Find the value at a key.
407 -- Returns 'Nothing' when the element can not be found.
408 --
409 -- > fromList [(5,'a'), (3,'b')] !? 1 == Nothing
410 -- > fromList [(5,'a'), (3,'b')] !? 5 == Just 'a'
411 --
412 -- @since 0.5.11
413
414 (!?) :: IntMap a -> Key -> Maybe a
415 (!?) m k = lookup k m
416
417 -- | Same as 'difference'.
418 (\\) :: IntMap a -> IntMap b -> IntMap a
419 m1 \\ m2 = difference m1 m2
420
421 infixl 9 !?,\\{-This comment teaches CPP correct behaviour -}
422
423 {--------------------------------------------------------------------
424 Types
425 --------------------------------------------------------------------}
426
427 instance Monoid (IntMap a) where
428 mempty = empty
429 mconcat = unions
430 #if !(MIN_VERSION_base(4,9,0))
431 mappend = union
432 #else
433 mappend = (<>)
434
435 -- | @since 0.5.7
436 instance Semigroup (IntMap a) where
437 (<>) = union
438 stimes = stimesIdempotentMonoid
439 #endif
440
441 instance Foldable.Foldable IntMap where
442 fold = go
443 where go Nil = mempty
444 go (Tip _ v) = v
445 go (Bin _ m l r)
446 | m < 0 = go r `mappend` go l
447 | otherwise = go l `mappend` go r
448 {-# INLINABLE fold #-}
449 foldr = foldr
450 {-# INLINE foldr #-}
451 foldl = foldl
452 {-# INLINE foldl #-}
453 foldMap f t = go t
454 where go Nil = mempty
455 go (Tip _ v) = f v
456 go (Bin _ m l r)
457 | m < 0 = go r `mappend` go l
458 | otherwise = go l `mappend` go r
459 {-# INLINE foldMap #-}
460 foldl' = foldl'
461 {-# INLINE foldl' #-}
462 foldr' = foldr'
463 {-# INLINE foldr' #-}
464 #if MIN_VERSION_base(4,8,0)
465 length = size
466 {-# INLINE length #-}
467 null = null
468 {-# INLINE null #-}
469 toList = elems -- NB: Foldable.toList /= IntMap.toList
470 {-# INLINE toList #-}
471 elem = go
472 where go !_ Nil = False
473 go x (Tip _ y) = x == y
474 go x (Bin _ _ l r) = go x l || go x r
475 {-# INLINABLE elem #-}
476 maximum = start
477 where start Nil = error "Data.Foldable.maximum (for Data.IntMap): empty map"
478 start (Tip _ y) = y
479 start (Bin _ _ l r) = go (start l) r
480
481 go !m Nil = m
482 go m (Tip _ y) = max m y
483 go m (Bin _ _ l r) = go (go m l) r
484 {-# INLINABLE maximum #-}
485 minimum = start
486 where start Nil = error "Data.Foldable.minimum (for Data.IntMap): empty map"
487 start (Tip _ y) = y
488 start (Bin _ _ l r) = go (start l) r
489
490 go !m Nil = m
491 go m (Tip _ y) = min m y
492 go m (Bin _ _ l r) = go (go m l) r
493 {-# INLINABLE minimum #-}
494 sum = foldl' (+) 0
495 {-# INLINABLE sum #-}
496 product = foldl' (*) 1
497 {-# INLINABLE product #-}
498 #endif
499
500 instance Traversable IntMap where
501 traverse f = traverseWithKey (\_ -> f)
502 {-# INLINE traverse #-}
503
504 instance NFData a => NFData (IntMap a) where
505 rnf Nil = ()
506 rnf (Tip _ v) = rnf v
507 rnf (Bin _ _ l r) = rnf l `seq` rnf r
508
509 #if __GLASGOW_HASKELL__
510
511 {--------------------------------------------------------------------
512 A Data instance
513 --------------------------------------------------------------------}
514
515 -- This instance preserves data abstraction at the cost of inefficiency.
516 -- We provide limited reflection services for the sake of data abstraction.
517
518 instance Data a => Data (IntMap a) where
519 gfoldl f z im = z fromList `f` (toList im)
520 toConstr _ = fromListConstr
521 gunfold k z c = case constrIndex c of
522 1 -> k (z fromList)
523 _ -> error "gunfold"
524 dataTypeOf _ = intMapDataType
525 dataCast1 f = gcast1 f
526
527 fromListConstr :: Constr
528 fromListConstr = mkConstr intMapDataType "fromList" [] Prefix
529
530 intMapDataType :: DataType
531 intMapDataType = mkDataType "Data.IntMap.Internal.IntMap" [fromListConstr]
532
533 #endif
534
535 {--------------------------------------------------------------------
536 Query
537 --------------------------------------------------------------------}
538 -- | /O(1)/. Is the map empty?
539 --
540 -- > Data.IntMap.null (empty) == True
541 -- > Data.IntMap.null (singleton 1 'a') == False
542
543 null :: IntMap a -> Bool
544 null Nil = True
545 null _ = False
546 {-# INLINE null #-}
547
548 -- | /O(n)/. Number of elements in the map.
549 --
550 -- > size empty == 0
551 -- > size (singleton 1 'a') == 1
552 -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3
553 size :: IntMap a -> Int
554 size = go 0
555 where
556 go !acc (Bin _ _ l r) = go (go acc l) r
557 go acc (Tip _ _) = 1 + acc
558 go acc Nil = acc
559
560 -- | /O(min(n,W))/. Is the key a member of the map?
561 --
562 -- > member 5 (fromList [(5,'a'), (3,'b')]) == True
563 -- > member 1 (fromList [(5,'a'), (3,'b')]) == False
564
565 -- See Note: Local 'go' functions and capturing]
566 member :: Key -> IntMap a -> Bool
567 member !k = go
568 where
569 go (Bin p m l r) | nomatch k p m = False
570 | zero k m = go l
571 | otherwise = go r
572 go (Tip kx _) = k == kx
573 go Nil = False
574
575 -- | /O(min(n,W))/. Is the key not a member of the map?
576 --
577 -- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False
578 -- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True
579
580 notMember :: Key -> IntMap a -> Bool
581 notMember k m = not $ member k m
582
583 -- | /O(min(n,W))/. Lookup the value at a key in the map. See also 'Data.Map.lookup'.
584
585 -- See Note: Local 'go' functions and capturing]
586 lookup :: Key -> IntMap a -> Maybe a
587 lookup !k = go
588 where
589 go (Bin p m l r) | nomatch k p m = Nothing
590 | zero k m = go l
591 | otherwise = go r
592 go (Tip kx x) | k == kx = Just x
593 | otherwise = Nothing
594 go Nil = Nothing
595
596
597 -- See Note: Local 'go' functions and capturing]
598 find :: Key -> IntMap a -> a
599 find !k = go
600 where
601 go (Bin p m l r) | nomatch k p m = not_found
602 | zero k m = go l
603 | otherwise = go r
604 go (Tip kx x) | k == kx = x
605 | otherwise = not_found
606 go Nil = not_found
607
608 not_found = error ("IntMap.!: key " ++ show k ++ " is not an element of the map")
609
610 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
611 -- returns the value at key @k@ or returns @def@ when the key is not an
612 -- element of the map.
613 --
614 -- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x'
615 -- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a'
616
617 -- See Note: Local 'go' functions and capturing]
618 findWithDefault :: a -> Key -> IntMap a -> a
619 findWithDefault def !k = go
620 where
621 go (Bin p m l r) | nomatch k p m = def
622 | zero k m = go l
623 | otherwise = go r
624 go (Tip kx x) | k == kx = x
625 | otherwise = def
626 go Nil = def
627
628 -- | /O(log n)/. Find largest key smaller than the given one and return the
629 -- corresponding (key, value) pair.
630 --
631 -- > lookupLT 3 (fromList [(3,'a'), (5,'b')]) == Nothing
632 -- > lookupLT 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
633
634 -- See Note: Local 'go' functions and capturing.
635 lookupLT :: Key -> IntMap a -> Maybe (Key, a)
636 lookupLT !k t = case t of
637 Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r
638 _ -> go Nil t
639 where
640 go def (Bin p m l r)
641 | nomatch k p m = if k < p then unsafeFindMax def else unsafeFindMax r
642 | zero k m = go def l
643 | otherwise = go l r
644 go def (Tip ky y)
645 | k <= ky = unsafeFindMax def
646 | otherwise = Just (ky, y)
647 go def Nil = unsafeFindMax def
648
649 -- | /O(log n)/. Find smallest key greater than the given one and return the
650 -- corresponding (key, value) pair.
651 --
652 -- > lookupGT 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
653 -- > lookupGT 5 (fromList [(3,'a'), (5,'b')]) == Nothing
654
655 -- See Note: Local 'go' functions and capturing.
656 lookupGT :: Key -> IntMap a -> Maybe (Key, a)
657 lookupGT !k t = case t of
658 Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r
659 _ -> go Nil t
660 where
661 go def (Bin p m l r)
662 | nomatch k p m = if k < p then unsafeFindMin l else unsafeFindMin def
663 | zero k m = go r l
664 | otherwise = go def r
665 go def (Tip ky y)
666 | k >= ky = unsafeFindMin def
667 | otherwise = Just (ky, y)
668 go def Nil = unsafeFindMin def
669
670 -- | /O(log n)/. Find largest key smaller or equal to the given one and return
671 -- the corresponding (key, value) pair.
672 --
673 -- > lookupLE 2 (fromList [(3,'a'), (5,'b')]) == Nothing
674 -- > lookupLE 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
675 -- > lookupLE 5 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
676
677 -- See Note: Local 'go' functions and capturing.
678 lookupLE :: Key -> IntMap a -> Maybe (Key, a)
679 lookupLE !k t = case t of
680 Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r
681 _ -> go Nil t
682 where
683 go def (Bin p m l r)
684 | nomatch k p m = if k < p then unsafeFindMax def else unsafeFindMax r
685 | zero k m = go def l
686 | otherwise = go l r
687 go def (Tip ky y)
688 | k < ky = unsafeFindMax def
689 | otherwise = Just (ky, y)
690 go def Nil = unsafeFindMax def
691
692 -- | /O(log n)/. Find smallest key greater or equal to the given one and return
693 -- the corresponding (key, value) pair.
694 --
695 -- > lookupGE 3 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
696 -- > lookupGE 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
697 -- > lookupGE 6 (fromList [(3,'a'), (5,'b')]) == Nothing
698
699 -- See Note: Local 'go' functions and capturing.
700 lookupGE :: Key -> IntMap a -> Maybe (Key, a)
701 lookupGE !k t = case t of
702 Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r
703 _ -> go Nil t
704 where
705 go def (Bin p m l r)
706 | nomatch k p m = if k < p then unsafeFindMin l else unsafeFindMin def
707 | zero k m = go r l
708 | otherwise = go def r
709 go def (Tip ky y)
710 | k > ky = unsafeFindMin def
711 | otherwise = Just (ky, y)
712 go def Nil = unsafeFindMin def
713
714
715 -- Helper function for lookupGE and lookupGT. It assumes that if a Bin node is
716 -- given, it has m > 0.
717 unsafeFindMin :: IntMap a -> Maybe (Key, a)
718 unsafeFindMin Nil = Nothing
719 unsafeFindMin (Tip ky y) = Just (ky, y)
720 unsafeFindMin (Bin _ _ l _) = unsafeFindMin l
721
722 -- Helper function for lookupLE and lookupLT. It assumes that if a Bin node is
723 -- given, it has m > 0.
724 unsafeFindMax :: IntMap a -> Maybe (Key, a)
725 unsafeFindMax Nil = Nothing
726 unsafeFindMax (Tip ky y) = Just (ky, y)
727 unsafeFindMax (Bin _ _ _ r) = unsafeFindMax r
728
729 {--------------------------------------------------------------------
730 Construction
731 --------------------------------------------------------------------}
732 -- | /O(1)/. The empty map.
733 --
734 -- > empty == fromList []
735 -- > size empty == 0
736
737 empty :: IntMap a
738 empty
739 = Nil
740 {-# INLINE empty #-}
741
742 -- | /O(1)/. A map of one element.
743 --
744 -- > singleton 1 'a' == fromList [(1, 'a')]
745 -- > size (singleton 1 'a') == 1
746
747 singleton :: Key -> a -> IntMap a
748 singleton k x
749 = Tip k x
750 {-# INLINE singleton #-}
751
752 {--------------------------------------------------------------------
753 Insert
754 --------------------------------------------------------------------}
755 -- | /O(min(n,W))/. Insert a new key\/value pair in the map.
756 -- If the key is already present in the map, the associated value is
757 -- replaced with the supplied value, i.e. 'insert' is equivalent to
758 -- @'insertWith' 'const'@.
759 --
760 -- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')]
761 -- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')]
762 -- > insert 5 'x' empty == singleton 5 'x'
763
764 insert :: Key -> a -> IntMap a -> IntMap a
765 insert !k x t@(Bin p m l r)
766 | nomatch k p m = link k (Tip k x) p t
767 | zero k m = Bin p m (insert k x l) r
768 | otherwise = Bin p m l (insert k x r)
769 insert k x t@(Tip ky _)
770 | k==ky = Tip k x
771 | otherwise = link k (Tip k x) ky t
772 insert k x Nil = Tip k x
773
774 -- right-biased insertion, used by 'union'
775 -- | /O(min(n,W))/. Insert with a combining function.
776 -- @'insertWith' f key value mp@
777 -- will insert the pair (key, value) into @mp@ if key does
778 -- not exist in the map. If the key does exist, the function will
779 -- insert @f new_value old_value@.
780 --
781 -- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")]
782 -- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
783 -- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx"
784
785 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
786 insertWith f k x t
787 = insertWithKey (\_ x' y' -> f x' y') k x t
788
789 -- | /O(min(n,W))/. Insert with a combining function.
790 -- @'insertWithKey' f key value mp@
791 -- will insert the pair (key, value) into @mp@ if key does
792 -- not exist in the map. If the key does exist, the function will
793 -- insert @f key new_value old_value@.
794 --
795 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
796 -- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")]
797 -- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
798 -- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx"
799
800 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
801 insertWithKey f !k x t@(Bin p m l r)
802 | nomatch k p m = link k (Tip k x) p t
803 | zero k m = Bin p m (insertWithKey f k x l) r
804 | otherwise = Bin p m l (insertWithKey f k x r)
805 insertWithKey f k x t@(Tip ky y)
806 | k == ky = Tip k (f k x y)
807 | otherwise = link k (Tip k x) ky t
808 insertWithKey _ k x Nil = Tip k x
809
810 -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
811 -- is a pair where the first element is equal to (@'lookup' k map@)
812 -- and the second element equal to (@'insertWithKey' f k x map@).
813 --
814 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
815 -- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")])
816 -- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")])
817 -- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx")
818 --
819 -- This is how to define @insertLookup@ using @insertLookupWithKey@:
820 --
821 -- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t
822 -- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")])
823 -- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")])
824
825 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
826 insertLookupWithKey f !k x t@(Bin p m l r)
827 | nomatch k p m = (Nothing,link k (Tip k x) p t)
828 | zero k m = let (found,l') = insertLookupWithKey f k x l
829 in (found,Bin p m l' r)
830 | otherwise = let (found,r') = insertLookupWithKey f k x r
831 in (found,Bin p m l r')
832 insertLookupWithKey f k x t@(Tip ky y)
833 | k == ky = (Just y,Tip k (f k x y))
834 | otherwise = (Nothing,link k (Tip k x) ky t)
835 insertLookupWithKey _ k x Nil = (Nothing,Tip k x)
836
837
838 {--------------------------------------------------------------------
839 Deletion
840 --------------------------------------------------------------------}
841 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
842 -- a member of the map, the original map is returned.
843 --
844 -- > delete 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
845 -- > delete 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
846 -- > delete 5 empty == empty
847
848 delete :: Key -> IntMap a -> IntMap a
849 delete !k t@(Bin p m l r)
850 | nomatch k p m = t
851 | zero k m = binCheckLeft p m (delete k l) r
852 | otherwise = binCheckRight p m l (delete k r)
853 delete k t@(Tip ky _)
854 | k == ky = Nil
855 | otherwise = t
856 delete _k Nil = Nil
857
858 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
859 -- a member of the map, the original map is returned.
860 --
861 -- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
862 -- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
863 -- > adjust ("new " ++) 7 empty == empty
864
865 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
866 adjust f k m
867 = adjustWithKey (\_ x -> f x) k m
868
869 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
870 -- a member of the map, the original map is returned.
871 --
872 -- > let f key x = (show key) ++ ":new " ++ x
873 -- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
874 -- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
875 -- > adjustWithKey f 7 empty == empty
876
877 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
878 adjustWithKey f !k t@(Bin p m l r)
879 | nomatch k p m = t
880 | zero k m = Bin p m (adjustWithKey f k l) r
881 | otherwise = Bin p m l (adjustWithKey f k r)
882 adjustWithKey f k t@(Tip ky y)
883 | k == ky = Tip ky (f k y)
884 | otherwise = t
885 adjustWithKey _ _ Nil = Nil
886
887
888 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
889 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
890 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
891 --
892 -- > let f x = if x == "a" then Just "new a" else Nothing
893 -- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
894 -- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
895 -- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
896
897 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
898 update f
899 = updateWithKey (\_ x -> f x)
900
901 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
902 -- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
903 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
904 --
905 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
906 -- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
907 -- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
908 -- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
909
910 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
911 updateWithKey f !k t@(Bin p m l r)
912 | nomatch k p m = t
913 | zero k m = binCheckLeft p m (updateWithKey f k l) r
914 | otherwise = binCheckRight p m l (updateWithKey f k r)
915 updateWithKey f k t@(Tip ky y)
916 | k == ky = case (f k y) of
917 Just y' -> Tip ky y'
918 Nothing -> Nil
919 | otherwise = t
920 updateWithKey _ _ Nil = Nil
921
922 -- | /O(min(n,W))/. Lookup and update.
923 -- The function returns original value, if it is updated.
924 -- This is different behavior than 'Data.Map.updateLookupWithKey'.
925 -- Returns the original key value if the map entry is deleted.
926 --
927 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
928 -- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:new a")])
929 -- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")])
930 -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a")
931
932 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
933 updateLookupWithKey f !k t@(Bin p m l r)
934 | nomatch k p m = (Nothing,t)
935 | zero k m = let !(found,l') = updateLookupWithKey f k l
936 in (found,binCheckLeft p m l' r)
937 | otherwise = let !(found,r') = updateLookupWithKey f k r
938 in (found,binCheckRight p m l r')
939 updateLookupWithKey f k t@(Tip ky y)
940 | k==ky = case (f k y) of
941 Just y' -> (Just y,Tip ky y')
942 Nothing -> (Just y,Nil)
943 | otherwise = (Nothing,t)
944 updateLookupWithKey _ _ Nil = (Nothing,Nil)
945
946
947
948 -- | /O(min(n,W))/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
949 -- 'alter' can be used to insert, delete, or update a value in an 'IntMap'.
950 -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
951 alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
952 alter f !k t@(Bin p m l r)
953 | nomatch k p m = case f Nothing of
954 Nothing -> t
955 Just x -> link k (Tip k x) p t
956 | zero k m = binCheckLeft p m (alter f k l) r
957 | otherwise = binCheckRight p m l (alter f k r)
958 alter f k t@(Tip ky y)
959 | k==ky = case f (Just y) of
960 Just x -> Tip ky x
961 Nothing -> Nil
962 | otherwise = case f Nothing of
963 Just x -> link k (Tip k x) ky t
964 Nothing -> Tip ky y
965 alter f k Nil = case f Nothing of
966 Just x -> Tip k x
967 Nothing -> Nil
968
969 -- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at
970 -- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete,
971 -- or update a value in an 'IntMap'. In short : @'lookup' k <$> 'alterF' f k m = f
972 -- ('lookup' k m)@.
973 --
974 -- Example:
975 --
976 -- @
977 -- interactiveAlter :: Int -> IntMap String -> IO (IntMap String)
978 -- interactiveAlter k m = alterF f k m where
979 -- f Nothing = do
980 -- putStrLn $ show k ++
981 -- " was not found in the map. Would you like to add it?"
982 -- getUserResponse1 :: IO (Maybe String)
983 -- f (Just old) = do
984 -- putStrLn $ "The key is currently bound to " ++ show old ++
985 -- ". Would you like to change or delete it?"
986 -- getUserResponse2 :: IO (Maybe String)
987 -- @
988 --
989 -- 'alterF' is the most general operation for working with an individual
990 -- key that may or may not be in a given map.
991 --
992 -- Note: 'alterF' is a flipped version of the @at@ combinator from
993 -- @Control.Lens.At@.
994 --
995 -- @since 0.5.8
996
997 alterF :: Functor f
998 => (Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a)
999 -- This implementation was stolen from 'Control.Lens.At'.
1000 alterF f k m = (<$> f mv) $ \fres ->
1001 case fres of
1002 Nothing -> maybe m (const (delete k m)) mv
1003 Just v' -> insert k v' m
1004 where mv = lookup k m
1005
1006 {--------------------------------------------------------------------
1007 Union
1008 --------------------------------------------------------------------}
1009 -- | The union of a list of maps.
1010 --
1011 -- > unions [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
1012 -- > == fromList [(3, "b"), (5, "a"), (7, "C")]
1013 -- > unions [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])]
1014 -- > == fromList [(3, "B3"), (5, "A3"), (7, "C")]
1015
1016 unions :: Foldable f => f (IntMap a) -> IntMap a
1017 unions xs
1018 = Foldable.foldl' union empty xs
1019
1020 -- | The union of a list of maps, with a combining operation.
1021 --
1022 -- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
1023 -- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")]
1024
1025 unionsWith :: Foldable f => (a->a->a) -> f (IntMap a) -> IntMap a
1026 unionsWith f ts
1027 = Foldable.foldl' (unionWith f) empty ts
1028
1029 -- | /O(n+m)/. The (left-biased) union of two maps.
1030 -- It prefers the first map when duplicate keys are encountered,
1031 -- i.e. (@'union' == 'unionWith' 'const'@).
1032 --
1033 -- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")]
1034
1035 union :: IntMap a -> IntMap a -> IntMap a
1036 union m1 m2
1037 = mergeWithKey' Bin const id id m1 m2
1038
1039 -- | /O(n+m)/. The union with a combining function.
1040 --
1041 -- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")]
1042
1043 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
1044 unionWith f m1 m2
1045 = unionWithKey (\_ x y -> f x y) m1 m2
1046
1047 -- | /O(n+m)/. The union with a combining function.
1048 --
1049 -- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value
1050 -- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")]
1051
1052 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
1053 unionWithKey f m1 m2
1054 = mergeWithKey' Bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 (f k1 x1 x2)) id id m1 m2
1055
1056 {--------------------------------------------------------------------
1057 Difference
1058 --------------------------------------------------------------------}
1059 -- | /O(n+m)/. Difference between two maps (based on keys).
1060 --
1061 -- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b"
1062
1063 difference :: IntMap a -> IntMap b -> IntMap a
1064 difference m1 m2
1065 = mergeWithKey (\_ _ _ -> Nothing) id (const Nil) m1 m2
1066
1067 -- | /O(n+m)/. Difference with a combining function.
1068 --
1069 -- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing
1070 -- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")])
1071 -- > == singleton 3 "b:B"
1072
1073 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
1074 differenceWith f m1 m2
1075 = differenceWithKey (\_ x y -> f x y) m1 m2
1076
1077 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
1078 -- encountered, the combining function is applied to the key and both values.
1079 -- If it returns 'Nothing', the element is discarded (proper set difference).
1080 -- If it returns (@'Just' y@), the element is updated with a new value @y@.
1081 --
1082 -- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing
1083 -- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")])
1084 -- > == singleton 3 "3:b|B"
1085
1086 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
1087 differenceWithKey f m1 m2
1088 = mergeWithKey f id (const Nil) m1 m2
1089
1090
1091 -- TODO(wrengr): re-verify that asymptotic bound
1092 -- | /O(n+m)/. Remove all the keys in a given set from a map.
1093 --
1094 -- @
1095 -- m \`withoutKeys\` s = 'filterWithKey' (\k _ -> k ``IntSet.notMember`` s) m
1096 -- @
1097 --
1098 -- @since 0.5.8
1099 withoutKeys :: IntMap a -> IntSet.IntSet -> IntMap a
1100 withoutKeys t1@(Bin p1 m1 l1 r1) t2@(IntSet.Bin p2 m2 l2 r2)
1101 | shorter m1 m2 = difference1
1102 | shorter m2 m1 = difference2
1103 | p1 == p2 = bin p1 m1 (withoutKeys l1 l2) (withoutKeys r1 r2)
1104 | otherwise = t1
1105 where
1106 difference1
1107 | nomatch p2 p1 m1 = t1
1108 | zero p2 m1 = binCheckLeft p1 m1 (withoutKeys l1 t2) r1
1109 | otherwise = binCheckRight p1 m1 l1 (withoutKeys r1 t2)
1110 difference2
1111 | nomatch p1 p2 m2 = t1
1112 | zero p1 m2 = withoutKeys t1 l2
1113 | otherwise = withoutKeys t1 r2
1114 withoutKeys t1@(Bin p1 m1 _ _) (IntSet.Tip p2 bm2) =
1115 let minbit = bitmapOf p1
1116 lt_minbit = minbit - 1
1117 maxbit = bitmapOf (p1 .|. (m1 .|. (m1 - 1)))
1118 gt_maxbit = maxbit `xor` complement (maxbit - 1)
1119 -- TODO(wrengr): should we manually inline/unroll 'updatePrefix'
1120 -- and 'withoutBM' here, in order to avoid redundant case analyses?
1121 in updatePrefix p2 t1 $ withoutBM (bm2 .|. lt_minbit .|. gt_maxbit)
1122 withoutKeys t1@(Bin _ _ _ _) IntSet.Nil = t1
1123 withoutKeys t1@(Tip k1 _) t2
1124 | k1 `IntSet.member` t2 = Nil
1125 | otherwise = t1
1126 withoutKeys Nil _ = Nil
1127
1128
1129 updatePrefix
1130 :: IntSetPrefix -> IntMap a -> (IntMap a -> IntMap a) -> IntMap a
1131 updatePrefix !kp t@(Bin p m l r) f
1132 | m .&. IntSet.suffixBitMask /= 0 =
1133 if p .&. IntSet.prefixBitMask == kp then f t else t
1134 | nomatch kp p m = t
1135 | zero kp m = binCheckLeft p m (updatePrefix kp l f) r
1136 | otherwise = binCheckRight p m l (updatePrefix kp r f)
1137 updatePrefix kp t@(Tip kx _) f
1138 | kx .&. IntSet.prefixBitMask == kp = f t
1139 | otherwise = t
1140 updatePrefix _ Nil _ = Nil
1141
1142
1143 withoutBM :: IntSetBitMap -> IntMap a -> IntMap a
1144 withoutBM 0 t = t
1145 withoutBM bm (Bin p m l r) =
1146 let leftBits = bitmapOf (p .|. m) - 1
1147 bmL = bm .&. leftBits
1148 bmR = bm `xor` bmL -- = (bm .&. complement leftBits)
1149 in bin p m (withoutBM bmL l) (withoutBM bmR r)
1150 withoutBM bm t@(Tip k _)
1151 -- TODO(wrengr): need we manually inline 'IntSet.Member' here?
1152 | k `IntSet.member` IntSet.Tip (k .&. IntSet.prefixBitMask) bm = Nil
1153 | otherwise = t
1154 withoutBM _ Nil = Nil
1155
1156
1157 {--------------------------------------------------------------------
1158 Intersection
1159 --------------------------------------------------------------------}
1160 -- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys).
1161 --
1162 -- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a"
1163
1164 intersection :: IntMap a -> IntMap b -> IntMap a
1165 intersection m1 m2
1166 = mergeWithKey' bin const (const Nil) (const Nil) m1 m2
1167
1168
1169 -- TODO(wrengr): re-verify that asymptotic bound
1170 -- | /O(n+m)/. The restriction of a map to the keys in a set.
1171 --
1172 -- @
1173 -- m \`restrictKeys\` s = 'filterWithKey' (\k _ -> k ``IntSet.member`` s) m
1174 -- @
1175 --
1176 -- @since 0.5.8
1177 restrictKeys :: IntMap a -> IntSet.IntSet -> IntMap a
1178 restrictKeys t1@(Bin p1 m1 l1 r1) t2@(IntSet.Bin p2 m2 l2 r2)
1179 | shorter m1 m2 = intersection1
1180 | shorter m2 m1 = intersection2
1181 | p1 == p2 = bin p1 m1 (restrictKeys l1 l2) (restrictKeys r1 r2)
1182 | otherwise = Nil
1183 where
1184 intersection1
1185 | nomatch p2 p1 m1 = Nil
1186 | zero p2 m1 = restrictKeys l1 t2
1187 | otherwise = restrictKeys r1 t2
1188 intersection2
1189 | nomatch p1 p2 m2 = Nil
1190 | zero p1 m2 = restrictKeys t1 l2
1191 | otherwise = restrictKeys t1 r2
1192 restrictKeys t1@(Bin p1 m1 _ _) (IntSet.Tip p2 bm2) =
1193 let minbit = bitmapOf p1
1194 ge_minbit = complement (minbit - 1)
1195 maxbit = bitmapOf (p1 .|. (m1 .|. (m1 - 1)))
1196 le_maxbit = maxbit .|. (maxbit - 1)
1197 -- TODO(wrengr): should we manually inline/unroll 'lookupPrefix'
1198 -- and 'restrictBM' here, in order to avoid redundant case analyses?
1199 in restrictBM (bm2 .&. ge_minbit .&. le_maxbit) (lookupPrefix p2 t1)
1200 restrictKeys (Bin _ _ _ _) IntSet.Nil = Nil
1201 restrictKeys t1@(Tip k1 _) t2
1202 | k1 `IntSet.member` t2 = t1
1203 | otherwise = Nil
1204 restrictKeys Nil _ = Nil
1205
1206
1207 -- | /O(min(n,W))/. Restrict to the sub-map with all keys matching
1208 -- a key prefix.
1209 lookupPrefix :: IntSetPrefix -> IntMap a -> IntMap a
1210 lookupPrefix !kp t@(Bin p m l r)
1211 | m .&. IntSet.suffixBitMask /= 0 =
1212 if p .&. IntSet.prefixBitMask == kp then t else Nil
1213 | nomatch kp p m = Nil
1214 | zero kp m = lookupPrefix kp l
1215 | otherwise = lookupPrefix kp r
1216 lookupPrefix kp t@(Tip kx _)
1217 | (kx .&. IntSet.prefixBitMask) == kp = t
1218 | otherwise = Nil
1219 lookupPrefix _ Nil = Nil
1220
1221
1222 restrictBM :: IntSetBitMap -> IntMap a -> IntMap a
1223 restrictBM 0 _ = Nil
1224 restrictBM bm (Bin p m l r) =
1225 let leftBits = bitmapOf (p .|. m) - 1
1226 bmL = bm .&. leftBits
1227 bmR = bm `xor` bmL -- = (bm .&. complement leftBits)
1228 in bin p m (restrictBM bmL l) (restrictBM bmR r)
1229 restrictBM bm t@(Tip k _)
1230 -- TODO(wrengr): need we manually inline 'IntSet.Member' here?
1231 | k `IntSet.member` IntSet.Tip (k .&. IntSet.prefixBitMask) bm = t
1232 | otherwise = Nil
1233 restrictBM _ Nil = Nil
1234
1235
1236 -- | /O(n+m)/. The intersection with a combining function.
1237 --
1238 -- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA"
1239
1240 intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
1241 intersectionWith f m1 m2
1242 = intersectionWithKey (\_ x y -> f x y) m1 m2
1243
1244 -- | /O(n+m)/. The intersection with a combining function.
1245 --
1246 -- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar
1247 -- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A"
1248
1249 intersectionWithKey :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
1250 intersectionWithKey f m1 m2
1251 = mergeWithKey' bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 (f k1 x1 x2)) (const Nil) (const Nil) m1 m2
1252
1253 {--------------------------------------------------------------------
1254 MergeWithKey
1255 --------------------------------------------------------------------}
1256
1257 -- | /O(n+m)/. A high-performance universal combining function. Using
1258 -- 'mergeWithKey', all combining functions can be defined without any loss of
1259 -- efficiency (with exception of 'union', 'difference' and 'intersection',
1260 -- where sharing of some nodes is lost with 'mergeWithKey').
1261 --
1262 -- Please make sure you know what is going on when using 'mergeWithKey',
1263 -- otherwise you can be surprised by unexpected code growth or even
1264 -- corruption of the data structure.
1265 --
1266 -- When 'mergeWithKey' is given three arguments, it is inlined to the call
1267 -- site. You should therefore use 'mergeWithKey' only to define your custom
1268 -- combining functions. For example, you could define 'unionWithKey',
1269 -- 'differenceWithKey' and 'intersectionWithKey' as
1270 --
1271 -- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2
1272 -- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2
1273 -- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2
1274 --
1275 -- When calling @'mergeWithKey' combine only1 only2@, a function combining two
1276 -- 'IntMap's is created, such that
1277 --
1278 -- * if a key is present in both maps, it is passed with both corresponding
1279 -- values to the @combine@ function. Depending on the result, the key is either
1280 -- present in the result with specified value, or is left out;
1281 --
1282 -- * a nonempty subtree present only in the first map is passed to @only1@ and
1283 -- the output is added to the result;
1284 --
1285 -- * a nonempty subtree present only in the second map is passed to @only2@ and
1286 -- the output is added to the result.
1287 --
1288 -- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/.
1289 -- The values can be modified arbitrarily. Most common variants of @only1@ and
1290 -- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@ or
1291 -- @'filterWithKey' f@ could be used for any @f@.
1292
1293 mergeWithKey :: (Key -> a -> b -> Maybe c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c)
1294 -> IntMap a -> IntMap b -> IntMap c
1295 mergeWithKey f g1 g2 = mergeWithKey' bin combine g1 g2
1296 where -- We use the lambda form to avoid non-exhaustive pattern matches warning.
1297 combine = \(Tip k1 x1) (Tip _k2 x2) ->
1298 case f k1 x1 x2 of
1299 Nothing -> Nil
1300 Just x -> Tip k1 x
1301 {-# INLINE combine #-}
1302 {-# INLINE mergeWithKey #-}
1303
1304 -- Slightly more general version of mergeWithKey. It differs in the following:
1305 --
1306 -- * the combining function operates on maps instead of keys and values. The
1307 -- reason is to enable sharing in union, difference and intersection.
1308 --
1309 -- * mergeWithKey' is given an equivalent of bin. The reason is that in union*,
1310 -- Bin constructor can be used, because we know both subtrees are nonempty.
1311
1312 mergeWithKey' :: (Prefix -> Mask -> IntMap c -> IntMap c -> IntMap c)
1313 -> (IntMap a -> IntMap b -> IntMap c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c)
1314 -> IntMap a -> IntMap b -> IntMap c
1315 mergeWithKey' bin' f g1 g2 = go
1316 where
1317 go t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
1318 | shorter m1 m2 = merge1
1319 | shorter m2 m1 = merge2
1320 | p1 == p2 = bin' p1 m1 (go l1 l2) (go r1 r2)
1321 | otherwise = maybe_link p1 (g1 t1) p2 (g2 t2)
1322 where
1323 merge1 | nomatch p2 p1 m1 = maybe_link p1 (g1 t1) p2 (g2 t2)
1324 | zero p2 m1 = bin' p1 m1 (go l1 t2) (g1 r1)
1325 | otherwise = bin' p1 m1 (g1 l1) (go r1 t2)
1326 merge2 | nomatch p1 p2 m2 = maybe_link p1 (g1 t1) p2 (g2 t2)
1327 | zero p1 m2 = bin' p2 m2 (go t1 l2) (g2 r2)
1328 | otherwise = bin' p2 m2 (g2 l2) (go t1 r2)
1329
1330 go t1'@(Bin _ _ _ _) t2'@(Tip k2' _) = merge0 t2' k2' t1'
1331 where
1332 merge0 t2 k2 t1@(Bin p1 m1 l1 r1)
1333 | nomatch k2 p1 m1 = maybe_link p1 (g1 t1) k2 (g2 t2)
1334 | zero k2 m1 = bin' p1 m1 (merge0 t2 k2 l1) (g1 r1)
1335 | otherwise = bin' p1 m1 (g1 l1) (merge0 t2 k2 r1)
1336 merge0 t2 k2 t1@(Tip k1 _)
1337 | k1 == k2 = f t1 t2
1338 | otherwise = maybe_link k1 (g1 t1) k2 (g2 t2)
1339 merge0 t2 _ Nil = g2 t2
1340
1341 go t1@(Bin _ _ _ _) Nil = g1 t1
1342
1343 go t1'@(Tip k1' _) t2' = merge0 t1' k1' t2'
1344 where
1345 merge0 t1 k1 t2@(Bin p2 m2 l2 r2)
1346 | nomatch k1 p2 m2 = maybe_link k1 (g1 t1) p2 (g2 t2)
1347 | zero k1 m2 = bin' p2 m2 (merge0 t1 k1 l2) (g2 r2)
1348 | otherwise = bin' p2 m2 (g2 l2) (merge0 t1 k1 r2)
1349 merge0 t1 k1 t2@(Tip k2 _)
1350 | k1 == k2 = f t1 t2
1351 | otherwise = maybe_link k1 (g1 t1) k2 (g2 t2)
1352 merge0 t1 _ Nil = g1 t1
1353
1354 go Nil t2 = g2 t2
1355
1356 maybe_link _ Nil _ t2 = t2
1357 maybe_link _ t1 _ Nil = t1
1358 maybe_link p1 t1 p2 t2 = link p1 t1 p2 t2
1359 {-# INLINE maybe_link #-}
1360 {-# INLINE mergeWithKey' #-}
1361
1362
1363 {--------------------------------------------------------------------
1364 mergeA
1365 --------------------------------------------------------------------}
1366
1367 -- | A tactic for dealing with keys present in one map but not the
1368 -- other in 'merge' or 'mergeA'.
1369 --
1370 -- A tactic of type @WhenMissing f k x z@ is an abstract representation
1371 -- of a function of type @Key -> x -> f (Maybe z)@.
1372 --
1373 -- @since 0.5.9
1374
1375 data WhenMissing f x y = WhenMissing
1376 { missingSubtree :: IntMap x -> f (IntMap y)
1377 , missingKey :: Key -> x -> f (Maybe y)}
1378
1379 -- | @since 0.5.9
1380 instance (Applicative f, Monad f) => Functor (WhenMissing f x) where
1381 fmap = mapWhenMissing
1382 {-# INLINE fmap #-}
1383
1384
1385 -- | @since 0.5.9
1386 instance (Applicative f, Monad f) => Category.Category (WhenMissing f)
1387 where
1388 id = preserveMissing
1389 f . g =
1390 traverseMaybeMissing $ \ k x -> do
1391 y <- missingKey g k x
1392 case y of
1393 Nothing -> pure Nothing
1394 Just q -> missingKey f k q
1395 {-# INLINE id #-}
1396 {-# INLINE (.) #-}
1397
1398
1399 -- | Equivalent to @ReaderT k (ReaderT x (MaybeT f))@.
1400 --
1401 -- @since 0.5.9
1402 instance (Applicative f, Monad f) => Applicative (WhenMissing f x) where
1403 pure x = mapMissing (\ _ _ -> x)
1404 f <*> g =
1405 traverseMaybeMissing $ \k x -> do
1406 res1 <- missingKey f k x
1407 case res1 of
1408 Nothing -> pure Nothing
1409 Just r -> (pure $!) . fmap r =<< missingKey g k x
1410 {-# INLINE pure #-}
1411 {-# INLINE (<*>) #-}
1412
1413
1414 -- | Equivalent to @ReaderT k (ReaderT x (MaybeT f))@.
1415 --
1416 -- @since 0.5.9
1417 instance (Applicative f, Monad f) => Monad (WhenMissing f x) where
1418 #if !MIN_VERSION_base(4,8,0)
1419 return = pure
1420 #endif
1421 m >>= f =
1422 traverseMaybeMissing $ \k x -> do
1423 res1 <- missingKey m k x
1424 case res1 of
1425 Nothing -> pure Nothing
1426 Just r -> missingKey (f r) k x
1427 {-# INLINE (>>=) #-}
1428
1429
1430 -- | Map covariantly over a @'WhenMissing' f x@.
1431 --
1432 -- @since 0.5.9
1433 mapWhenMissing
1434 :: (Applicative f, Monad f)
1435 => (a -> b)
1436 -> WhenMissing f x a
1437 -> WhenMissing f x b
1438 mapWhenMissing f t = WhenMissing
1439 { missingSubtree = \m -> missingSubtree t m >>= \m' -> pure $! fmap f m'
1440 , missingKey = \k x -> missingKey t k x >>= \q -> (pure $! fmap f q) }
1441 {-# INLINE mapWhenMissing #-}
1442
1443
1444 -- | Map covariantly over a @'WhenMissing' f x@, using only a
1445 -- 'Functor f' constraint.
1446 mapGentlyWhenMissing
1447 :: Functor f
1448 => (a -> b)
1449 -> WhenMissing f x a
1450 -> WhenMissing f x b
1451 mapGentlyWhenMissing f t = WhenMissing
1452 { missingSubtree = \m -> fmap f <$> missingSubtree t m
1453 , missingKey = \k x -> fmap f <$> missingKey t k x }
1454 {-# INLINE mapGentlyWhenMissing #-}
1455
1456
1457 -- | Map covariantly over a @'WhenMatched' f k x@, using only a
1458 -- 'Functor f' constraint.
1459 mapGentlyWhenMatched
1460 :: Functor f
1461 => (a -> b)
1462 -> WhenMatched f x y a
1463 -> WhenMatched f x y b
1464 mapGentlyWhenMatched f t =
1465 zipWithMaybeAMatched $ \k x y -> fmap f <$> runWhenMatched t k x y
1466 {-# INLINE mapGentlyWhenMatched #-}
1467
1468
1469 -- | Map contravariantly over a @'WhenMissing' f _ x@.
1470 --
1471 -- @since 0.5.9
1472 lmapWhenMissing :: (b -> a) -> WhenMissing f a x -> WhenMissing f b x
1473 lmapWhenMissing f t = WhenMissing
1474 { missingSubtree = \m -> missingSubtree t (fmap f m)
1475 , missingKey = \k x -> missingKey t k (f x) }
1476 {-# INLINE lmapWhenMissing #-}
1477
1478
1479 -- | Map contravariantly over a @'WhenMatched' f _ y z@.
1480 --
1481 -- @since 0.5.9
1482 contramapFirstWhenMatched
1483 :: (b -> a)
1484 -> WhenMatched f a y z
1485 -> WhenMatched f b y z
1486 contramapFirstWhenMatched f t =
1487 WhenMatched $ \k x y -> runWhenMatched t k (f x) y
1488 {-# INLINE contramapFirstWhenMatched #-}
1489
1490
1491 -- | Map contravariantly over a @'WhenMatched' f x _ z@.
1492 --
1493 -- @since 0.5.9
1494 contramapSecondWhenMatched
1495 :: (b -> a)
1496 -> WhenMatched f x a z
1497 -> WhenMatched f x b z
1498 contramapSecondWhenMatched f t =
1499 WhenMatched $ \k x y -> runWhenMatched t k x (f y)
1500 {-# INLINE contramapSecondWhenMatched #-}
1501
1502
1503 #if !MIN_VERSION_base(4,8,0)
1504 newtype Identity a = Identity {runIdentity :: a}
1505
1506 instance Functor Identity where
1507 fmap f (Identity x) = Identity (f x)
1508
1509 instance Applicative Identity where
1510 pure = Identity
1511 Identity f <*> Identity x = Identity (f x)
1512 #endif
1513
1514 -- | A tactic for dealing with keys present in one map but not the
1515 -- other in 'merge'.
1516 --
1517 -- A tactic of type @SimpleWhenMissing x z@ is an abstract
1518 -- representation of a function of type @Key -> x -> Maybe z@.
1519 --
1520 -- @since 0.5.9
1521 type SimpleWhenMissing = WhenMissing Identity
1522
1523
1524 -- | A tactic for dealing with keys present in both maps in 'merge'
1525 -- or 'mergeA'.
1526 --
1527 -- A tactic of type @WhenMatched f x y z@ is an abstract representation
1528 -- of a function of type @Key -> x -> y -> f (Maybe z)@.
1529 --
1530 -- @since 0.5.9
1531 newtype WhenMatched f x y z = WhenMatched
1532 { matchedKey :: Key -> x -> y -> f (Maybe z) }
1533
1534
1535 -- | Along with zipWithMaybeAMatched, witnesses the isomorphism
1536 -- between @WhenMatched f x y z@ and @Key -> x -> y -> f (Maybe z)@.
1537 --
1538 -- @since 0.5.9
1539 runWhenMatched :: WhenMatched f x y z -> Key -> x -> y -> f (Maybe z)
1540 runWhenMatched = matchedKey
1541 {-# INLINE runWhenMatched #-}
1542
1543
1544 -- | Along with traverseMaybeMissing, witnesses the isomorphism
1545 -- between @WhenMissing f x y@ and @Key -> x -> f (Maybe y)@.
1546 --
1547 -- @since 0.5.9
1548 runWhenMissing :: WhenMissing f x y -> Key-> x -> f (Maybe y)
1549 runWhenMissing = missingKey
1550 {-# INLINE runWhenMissing #-}
1551
1552
1553 -- | @since 0.5.9
1554 instance Functor f => Functor (WhenMatched f x y) where
1555 fmap = mapWhenMatched
1556 {-# INLINE fmap #-}
1557
1558
1559 -- | @since 0.5.9
1560 instance (Monad f, Applicative f) => Category.Category (WhenMatched f x)
1561 where
1562 id = zipWithMatched (\_ _ y -> y)
1563 f . g =
1564 zipWithMaybeAMatched $ \k x y -> do
1565 res <- runWhenMatched g k x y
1566 case res of
1567 Nothing -> pure Nothing
1568 Just r -> runWhenMatched f k x r
1569 {-# INLINE id #-}
1570 {-# INLINE (.) #-}
1571
1572
1573 -- | Equivalent to @ReaderT Key (ReaderT x (ReaderT y (MaybeT f)))@
1574 --
1575 -- @since 0.5.9
1576 instance (Monad f, Applicative f) => Applicative (WhenMatched f x y) where
1577 pure x = zipWithMatched (\_ _ _ -> x)
1578 fs <*> xs =
1579 zipWithMaybeAMatched $ \k x y -> do
1580 res <- runWhenMatched fs k x y
1581 case res of
1582 Nothing -> pure Nothing
1583 Just r -> (pure $!) . fmap r =<< runWhenMatched xs k x y
1584 {-# INLINE pure #-}
1585 {-# INLINE (<*>) #-}
1586
1587
1588 -- | Equivalent to @ReaderT Key (ReaderT x (ReaderT y (MaybeT f)))@
1589 --
1590 -- @since 0.5.9
1591 instance (Monad f, Applicative f) => Monad (WhenMatched f x y) where
1592 #if !MIN_VERSION_base(4,8,0)
1593 return = pure
1594 #endif
1595 m >>= f =
1596 zipWithMaybeAMatched $ \k x y -> do
1597 res <- runWhenMatched m k x y
1598 case res of
1599 Nothing -> pure Nothing
1600 Just r -> runWhenMatched (f r) k x y
1601 {-# INLINE (>>=) #-}
1602
1603
1604 -- | Map covariantly over a @'WhenMatched' f x y@.
1605 --
1606 -- @since 0.5.9
1607 mapWhenMatched
1608 :: Functor f
1609 => (a -> b)
1610 -> WhenMatched f x y a
1611 -> WhenMatched f x y b
1612 mapWhenMatched f (WhenMatched g) =
1613 WhenMatched $ \k x y -> fmap (fmap f) (g k x y)
1614 {-# INLINE mapWhenMatched #-}
1615
1616
1617 -- | A tactic for dealing with keys present in both maps in 'merge'.
1618 --
1619 -- A tactic of type @SimpleWhenMatched x y z@ is an abstract
1620 -- representation of a function of type @Key -> x -> y -> Maybe z@.
1621 --
1622 -- @since 0.5.9
1623 type SimpleWhenMatched = WhenMatched Identity
1624
1625
1626 -- | When a key is found in both maps, apply a function to the key
1627 -- and values and use the result in the merged map.
1628 --
1629 -- > zipWithMatched
1630 -- > :: (Key -> x -> y -> z)
1631 -- > -> SimpleWhenMatched x y z
1632 --
1633 -- @since 0.5.9
1634 zipWithMatched
1635 :: Applicative f
1636 => (Key -> x -> y -> z)
1637 -> WhenMatched f x y z
1638 zipWithMatched f = WhenMatched $ \ k x y -> pure . Just $ f k x y
1639 {-# INLINE zipWithMatched #-}
1640
1641
1642 -- | When a key is found in both maps, apply a function to the key
1643 -- and values to produce an action and use its result in the merged
1644 -- map.
1645 --
1646 -- @since 0.5.9
1647 zipWithAMatched
1648 :: Applicative f
1649 => (Key -> x -> y -> f z)
1650 -> WhenMatched f x y z
1651 zipWithAMatched f = WhenMatched $ \ k x y -> Just <$> f k x y
1652 {-# INLINE zipWithAMatched #-}
1653
1654
1655 -- | When a key is found in both maps, apply a function to the key
1656 -- and values and maybe use the result in the merged map.
1657 --
1658 -- > zipWithMaybeMatched
1659 -- > :: (Key -> x -> y -> Maybe z)
1660 -- > -> SimpleWhenMatched x y z
1661 --
1662 -- @since 0.5.9
1663 zipWithMaybeMatched
1664 :: Applicative f
1665 => (Key -> x -> y -> Maybe z)
1666 -> WhenMatched f x y z
1667 zipWithMaybeMatched f = WhenMatched $ \ k x y -> pure $ f k x y
1668 {-# INLINE zipWithMaybeMatched #-}
1669
1670
1671 -- | When a key is found in both maps, apply a function to the key
1672 -- and values, perform the resulting action, and maybe use the
1673 -- result in the merged map.
1674 --
1675 -- This is the fundamental 'WhenMatched' tactic.
1676 --
1677 -- @since 0.5.9
1678 zipWithMaybeAMatched
1679 :: (Key -> x -> y -> f (Maybe z))
1680 -> WhenMatched f x y z
1681 zipWithMaybeAMatched f = WhenMatched $ \ k x y -> f k x y
1682 {-# INLINE zipWithMaybeAMatched #-}
1683
1684
1685 -- | Drop all the entries whose keys are missing from the other
1686 -- map.
1687 --
1688 -- > dropMissing :: SimpleWhenMissing x y
1689 --
1690 -- prop> dropMissing = mapMaybeMissing (\_ _ -> Nothing)
1691 --
1692 -- but @dropMissing@ is much faster.
1693 --
1694 -- @since 0.5.9
1695 dropMissing :: Applicative f => WhenMissing f x y
1696 dropMissing = WhenMissing
1697 { missingSubtree = const (pure Nil)
1698 , missingKey = \_ _ -> pure Nothing }
1699 {-# INLINE dropMissing #-}
1700
1701
1702 -- | Preserve, unchanged, the entries whose keys are missing from
1703 -- the other map.
1704 --
1705 -- > preserveMissing :: SimpleWhenMissing x x
1706 --
1707 -- prop> preserveMissing = Merge.Lazy.mapMaybeMissing (\_ x -> Just x)
1708 --
1709 -- but @preserveMissing@ is much faster.
1710 --
1711 -- @since 0.5.9
1712 preserveMissing :: Applicative f => WhenMissing f x x
1713 preserveMissing = WhenMissing
1714 { missingSubtree = pure
1715 , missingKey = \_ v -> pure (Just v) }
1716 {-# INLINE preserveMissing #-}
1717
1718
1719 -- | Map over the entries whose keys are missing from the other map.
1720 --
1721 -- > mapMissing :: (k -> x -> y) -> SimpleWhenMissing x y
1722 --
1723 -- prop> mapMissing f = mapMaybeMissing (\k x -> Just $ f k x)
1724 --
1725 -- but @mapMissing@ is somewhat faster.
1726 --
1727 -- @since 0.5.9
1728 mapMissing :: Applicative f => (Key -> x -> y) -> WhenMissing f x y
1729 mapMissing f = WhenMissing
1730 { missingSubtree = \m -> pure $! mapWithKey f m
1731 , missingKey = \k x -> pure $ Just (f k x) }
1732 {-# INLINE mapMissing #-}
1733
1734
1735 -- | Map over the entries whose keys are missing from the other
1736 -- map, optionally removing some. This is the most powerful
1737 -- 'SimpleWhenMissing' tactic, but others are usually more efficient.
1738 --
1739 -- > mapMaybeMissing :: (Key -> x -> Maybe y) -> SimpleWhenMissing x y
1740 --
1741 -- prop> mapMaybeMissing f = traverseMaybeMissing (\k x -> pure (f k x))
1742 --
1743 -- but @mapMaybeMissing@ uses fewer unnecessary 'Applicative'
1744 -- operations.
1745 --
1746 -- @since 0.5.9
1747 mapMaybeMissing
1748 :: Applicative f => (Key -> x -> Maybe y) -> WhenMissing f x y
1749 mapMaybeMissing f = WhenMissing
1750 { missingSubtree = \m -> pure $! mapMaybeWithKey f m
1751 , missingKey = \k x -> pure $! f k x }
1752 {-# INLINE mapMaybeMissing #-}
1753
1754
1755 -- | Filter the entries whose keys are missing from the other map.
1756 --
1757 -- > filterMissing :: (k -> x -> Bool) -> SimpleWhenMissing x x
1758 --
1759 -- prop> filterMissing f = Merge.Lazy.mapMaybeMissing $ \k x -> guard (f k x) *> Just x
1760 --
1761 -- but this should be a little faster.
1762 --
1763 -- @since 0.5.9
1764 filterMissing
1765 :: Applicative f => (Key -> x -> Bool) -> WhenMissing f x x
1766 filterMissing f = WhenMissing
1767 { missingSubtree = \m -> pure $! filterWithKey f m
1768 , missingKey = \k x -> pure $! if f k x then Just x else Nothing }
1769 {-# INLINE filterMissing #-}
1770
1771
1772 -- | Filter the entries whose keys are missing from the other map
1773 -- using some 'Applicative' action.
1774 --
1775 -- > filterAMissing f = Merge.Lazy.traverseMaybeMissing $
1776 -- > \k x -> (\b -> guard b *> Just x) <$> f k x
1777 --
1778 -- but this should be a little faster.
1779 --
1780 -- @since 0.5.9
1781 filterAMissing
1782 :: Applicative f => (Key -> x -> f Bool) -> WhenMissing f x x
1783 filterAMissing f = WhenMissing
1784 { missingSubtree = \m -> filterWithKeyA f m
1785 , missingKey = \k x -> bool Nothing (Just x) <$> f k x }
1786 {-# INLINE filterAMissing #-}
1787
1788
1789 -- | /O(n)/. Filter keys and values using an 'Applicative' predicate.
1790 filterWithKeyA
1791 :: Applicative f => (Key -> a -> f Bool) -> IntMap a -> f (IntMap a)
1792 filterWithKeyA _ Nil = pure Nil
1793 filterWithKeyA f t@(Tip k x) = (\b -> if b then t else Nil) <$> f k x
1794 filterWithKeyA f (Bin p m l r) =
1795 liftA2 (bin p m) (filterWithKeyA f l) (filterWithKeyA f r)
1796
1797 -- | This wasn't in Data.Bool until 4.7.0, so we define it here
1798 bool :: a -> a -> Bool -> a
1799 bool f _ False = f
1800 bool _ t True = t
1801
1802
1803 -- | Traverse over the entries whose keys are missing from the other
1804 -- map.
1805 --
1806 -- @since 0.5.9
1807 traverseMissing
1808 :: Applicative f => (Key -> x -> f y) -> WhenMissing f x y
1809 traverseMissing f = WhenMissing
1810 { missingSubtree = traverseWithKey f
1811 , missingKey = \k x -> Just <$> f k x }
1812 {-# INLINE traverseMissing #-}
1813
1814
1815 -- | Traverse over the entries whose keys are missing from the other
1816 -- map, optionally producing values to put in the result. This is
1817 -- the most powerful 'WhenMissing' tactic, but others are usually
1818 -- more efficient.
1819 --
1820 -- @since 0.5.9
1821 traverseMaybeMissing
1822 :: Applicative f => (Key -> x -> f (Maybe y)) -> WhenMissing f x y
1823 traverseMaybeMissing f = WhenMissing
1824 { missingSubtree = traverseMaybeWithKey f
1825 , missingKey = f }
1826 {-# INLINE traverseMaybeMissing #-}
1827
1828
1829 -- | /O(n)/. Traverse keys\/values and collect the 'Just' results.
1830 traverseMaybeWithKey
1831 :: Applicative f => (Key -> a -> f (Maybe b)) -> IntMap a -> f (IntMap b)
1832 traverseMaybeWithKey f = go
1833 where
1834 go Nil = pure Nil
1835 go (Tip k x) = maybe Nil (Tip k) <$> f k x
1836 go (Bin p m l r) = liftA2 (bin p m) (go l) (go r)
1837
1838
1839 -- | Merge two maps.
1840 --
1841 -- 'merge' takes two 'WhenMissing' tactics, a 'WhenMatched' tactic
1842 -- and two maps. It uses the tactics to merge the maps. Its behavior
1843 -- is best understood via its fundamental tactics, 'mapMaybeMissing'
1844 -- and 'zipWithMaybeMatched'.
1845 --
1846 -- Consider
1847 --
1848 -- @
1849 -- merge (mapMaybeMissing g1)
1850 -- (mapMaybeMissing g2)
1851 -- (zipWithMaybeMatched f)
1852 -- m1 m2
1853 -- @
1854 --
1855 -- Take, for example,
1856 --
1857 -- @
1858 -- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')]
1859 -- m2 = [(1, "one"), (2, "two"), (4, "three")]
1860 -- @
1861 --
1862 -- 'merge' will first \"align\" these maps by key:
1863 --
1864 -- @
1865 -- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')]
1866 -- m2 = [(1, "one"), (2, "two"), (4, "three")]
1867 -- @
1868 --
1869 -- It will then pass the individual entries and pairs of entries
1870 -- to @g1@, @g2@, or @f@ as appropriate:
1871 --
1872 -- @
1873 -- maybes = [g1 0 \'a\', f 1 \'b\' "one", g2 2 "two", g1 3 \'c\', f 4 \'d\' "three"]
1874 -- @
1875 --
1876 -- This produces a 'Maybe' for each key:
1877 --
1878 -- @
1879 -- keys = 0 1 2 3 4
1880 -- results = [Nothing, Just True, Just False, Nothing, Just True]
1881 -- @
1882 --
1883 -- Finally, the @Just@ results are collected into a map:
1884 --
1885 -- @
1886 -- return value = [(1, True), (2, False), (4, True)]
1887 -- @
1888 --
1889 -- The other tactics below are optimizations or simplifications of
1890 -- 'mapMaybeMissing' for special cases. Most importantly,
1891 --
1892 -- * 'dropMissing' drops all the keys.
1893 -- * 'preserveMissing' leaves all the entries alone.
1894 --
1895 -- When 'merge' is given three arguments, it is inlined at the call
1896 -- site. To prevent excessive inlining, you should typically use
1897 -- 'merge' to define your custom combining functions.
1898 --
1899 --
1900 -- Examples:
1901 --
1902 -- prop> unionWithKey f = merge preserveMissing preserveMissing (zipWithMatched f)
1903 -- prop> intersectionWithKey f = merge dropMissing dropMissing (zipWithMatched f)
1904 -- prop> differenceWith f = merge diffPreserve diffDrop f
1905 -- prop> symmetricDifference = merge diffPreserve diffPreserve (\ _ _ _ -> Nothing)
1906 -- prop> mapEachPiece f g h = merge (diffMapWithKey f) (diffMapWithKey g)
1907 --
1908 -- @since 0.5.9
1909 merge
1910 :: SimpleWhenMissing a c -- ^ What to do with keys in @m1@ but not @m2@
1911 -> SimpleWhenMissing b c -- ^ What to do with keys in @m2@ but not @m1@
1912 -> SimpleWhenMatched a b c -- ^ What to do with keys in both @m1@ and @m2@
1913 -> IntMap a -- ^ Map @m1@
1914 -> IntMap b -- ^ Map @m2@
1915 -> IntMap c
1916 merge g1 g2 f m1 m2 =
1917 runIdentity $ mergeA g1 g2 f m1 m2
1918 {-# INLINE merge #-}
1919
1920
1921 -- | An applicative version of 'merge'.
1922 --
1923 -- 'mergeA' takes two 'WhenMissing' tactics, a 'WhenMatched'
1924 -- tactic and two maps. It uses the tactics to merge the maps.
1925 -- Its behavior is best understood via its fundamental tactics,
1926 -- 'traverseMaybeMissing' and 'zipWithMaybeAMatched'.
1927 --
1928 -- Consider
1929 --
1930 -- @
1931 -- mergeA (traverseMaybeMissing g1)
1932 -- (traverseMaybeMissing g2)
1933 -- (zipWithMaybeAMatched f)
1934 -- m1 m2
1935 -- @
1936 --
1937 -- Take, for example,
1938 --
1939 -- @
1940 -- m1 = [(0, \'a\'), (1, \'b\'), (3,\'c\'), (4, \'d\')]
1941 -- m2 = [(1, "one"), (2, "two"), (4, "three")]
1942 -- @
1943 --
1944 -- 'mergeA' will first \"align\" these maps by key:
1945 --
1946 -- @
1947 -- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')]
1948 -- m2 = [(1, "one"), (2, "two"), (4, "three")]
1949 -- @
1950 --
1951 -- It will then pass the individual entries and pairs of entries
1952 -- to @g1@, @g2@, or @f@ as appropriate:
1953 --
1954 -- @
1955 -- actions = [g1 0 \'a\', f 1 \'b\' "one", g2 2 "two", g1 3 \'c\', f 4 \'d\' "three"]
1956 -- @
1957 --
1958 -- Next, it will perform the actions in the @actions@ list in order from
1959 -- left to right.
1960 --
1961 -- @
1962 -- keys = 0 1 2 3 4
1963 -- results = [Nothing, Just True, Just False, Nothing, Just True]
1964 -- @
1965 --
1966 -- Finally, the @Just@ results are collected into a map:
1967 --
1968 -- @
1969 -- return value = [(1, True), (2, False), (4, True)]
1970 -- @
1971 --
1972 -- The other tactics below are optimizations or simplifications of
1973 -- 'traverseMaybeMissing' for special cases. Most importantly,
1974 --
1975 -- * 'dropMissing' drops all the keys.
1976 -- * 'preserveMissing' leaves all the entries alone.
1977 -- * 'mapMaybeMissing' does not use the 'Applicative' context.
1978 --
1979 -- When 'mergeA' is given three arguments, it is inlined at the call
1980 -- site. To prevent excessive inlining, you should generally only use
1981 -- 'mergeA' to define custom combining functions.
1982 --
1983 -- @since 0.5.9
1984 mergeA
1985 :: (Applicative f)
1986 => WhenMissing f a c -- ^ What to do with keys in @m1@ but not @m2@
1987 -> WhenMissing f b c -- ^ What to do with keys in @m2@ but not @m1@
1988 -> WhenMatched f a b c -- ^ What to do with keys in both @m1@ and @m2@
1989 -> IntMap a -- ^ Map @m1@
1990 -> IntMap b -- ^ Map @m2@
1991 -> f (IntMap c)
1992 mergeA
1993 WhenMissing{missingSubtree = g1t, missingKey = g1k}
1994 WhenMissing{missingSubtree = g2t, missingKey = g2k}
1995 WhenMatched{matchedKey = f}
1996 = go
1997 where
1998 go t1 Nil = g1t t1
1999 go Nil t2 = g2t t2
2000
2001 -- This case is already covered below.
2002 -- go (Tip k1 x1) (Tip k2 x2) = mergeTips k1 x1 k2 x2
2003
2004 go (Tip k1 x1) t2' = merge2 t2'
2005 where
2006 merge2 t2@(Bin p2 m2 l2 r2)
2007 | nomatch k1 p2 m2 = linkA k1 (subsingletonBy g1k k1 x1) p2 (g2t t2)
2008 | zero k1 m2 = liftA2 (bin p2 m2) (merge2 l2) (g2t r2)
2009 | otherwise = liftA2 (bin p2 m2) (g2t l2) (merge2 r2)
2010 merge2 (Tip k2 x2) = mergeTips k1 x1 k2 x2
2011 merge2 Nil = subsingletonBy g1k k1 x1
2012
2013 go t1' (Tip k2 x2) = merge1 t1'
2014 where
2015 merge1 t1@(Bin p1 m1 l1 r1)
2016 | nomatch k2 p1 m1 = linkA p1 (g1t t1) k2 (subsingletonBy g2k k2 x2)
2017 | zero k2 m1 = liftA2 (bin p1 m1) (merge1 l1) (g1t r1)
2018 | otherwise = liftA2 (bin p1 m1) (g1t l1) (merge1 r1)
2019 merge1 (Tip k1 x1) = mergeTips k1 x1 k2 x2
2020 merge1 Nil = subsingletonBy g2k k2 x2
2021
2022 go t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
2023 | shorter m1 m2 = merge1
2024 | shorter m2 m1 = merge2
2025 | p1 == p2 = liftA2 (bin p1 m1) (go l1 l2) (go r1 r2)
2026 | otherwise = liftA2 (link_ p1 p2) (g1t t1) (g2t t2)
2027 where
2028 merge1 | nomatch p2 p1 m1 = liftA2 (link_ p1 p2) (g1t t1) (g2t t2)
2029 | zero p2 m1 = liftA2 (bin p1 m1) (go l1 t2) (g1t r1)
2030 | otherwise = liftA2 (bin p1 m1) (g1t l1) (go r1 t2)
2031 merge2 | nomatch p1 p2 m2 = liftA2 (link_ p1 p2) (g1t t1) (g2t t2)
2032 | zero p1 m2 = liftA2 (bin p2 m2) (go t1 l2) (g2t r2)
2033 | otherwise = liftA2 (bin p2 m2) (g2t l2) (go t1 r2)
2034
2035 subsingletonBy gk k x = maybe Nil (Tip k) <$> gk k x
2036 {-# INLINE subsingletonBy #-}
2037
2038 mergeTips k1 x1 k2 x2
2039 | k1 == k2 = maybe Nil (Tip k1) <$> f k1 x1 x2
2040 | k1 < k2 = liftA2 (subdoubleton k1 k2) (g1k k1 x1) (g2k k2 x2)
2041 {-
2042 = link_ k1 k2 <$> subsingletonBy g1k k1 x1 <*> subsingletonBy g2k k2 x2
2043 -}
2044 | otherwise = liftA2 (subdoubleton k2 k1) (g2k k2 x2) (g1k k1 x1)
2045 {-# INLINE mergeTips #-}
2046
2047 subdoubleton _ _ Nothing Nothing = Nil
2048 subdoubleton _ k2 Nothing (Just y2) = Tip k2 y2
2049 subdoubleton k1 _ (Just y1) Nothing = Tip k1 y1
2050 subdoubleton k1 k2 (Just y1) (Just y2) = link k1 (Tip k1 y1) k2 (Tip k2 y2)
2051 {-# INLINE subdoubleton #-}
2052
2053 link_ _ _ Nil t2 = t2
2054 link_ _ _ t1 Nil = t1
2055 link_ p1 p2 t1 t2 = link p1 t1 p2 t2
2056 {-# INLINE link_ #-}
2057
2058 -- | A variant of 'link_' which makes sure to execute side-effects
2059 -- in the right order.
2060 linkA
2061 :: Applicative f
2062 => Prefix -> f (IntMap a)
2063 -> Prefix -> f (IntMap a)
2064 -> f (IntMap a)
2065 linkA p1 t1 p2 t2
2066 | zero p1 m = liftA2 (bin p m) t1 t2
2067 | otherwise = liftA2 (bin p m) t2 t1
2068 where
2069 m = branchMask p1 p2
2070 p = mask p1 m
2071 {-# INLINE linkA #-}
2072 {-# INLINE mergeA #-}
2073
2074
2075 {--------------------------------------------------------------------
2076 Min\/Max
2077 --------------------------------------------------------------------}
2078
2079 -- | /O(min(n,W))/. Update the value at the minimal key.
2080 --
2081 -- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"3:b"), (5,"a")]
2082 -- > updateMinWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
2083
2084 updateMinWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a
2085 updateMinWithKey f t =
2086 case t of Bin p m l r | m < 0 -> binCheckRight p m l (go f r)
2087 _ -> go f t
2088 where
2089 go f' (Bin p m l r) = binCheckLeft p m (go f' l) r
2090 go f' (Tip k y) = case f' k y of
2091 Just y' -> Tip k y'
2092 Nothing -> Nil
2093 go _ Nil = error "updateMinWithKey Nil"
2094
2095 -- | /O(min(n,W))/. Update the value at the maximal key.
2096 --
2097 -- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")]
2098 -- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
2099
2100 updateMaxWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a
2101 updateMaxWithKey f t =
2102 case t of Bin p m l r | m < 0 -> binCheckLeft p m (go f l) r
2103 _ -> go f t
2104 where
2105 go f' (Bin p m l r) = binCheckRight p m l (go f' r)
2106 go f' (Tip k y) = case f' k y of
2107 Just y' -> Tip k y'
2108 Nothing -> Nil
2109 go _ Nil = error "updateMaxWithKey Nil"
2110
2111
2112 data View a = View {-# UNPACK #-} !Key a !(IntMap a)
2113
2114 -- | /O(min(n,W))/. Retrieves the maximal (key,value) pair of the map, and
2115 -- the map stripped of that element, or 'Nothing' if passed an empty map.
2116 --
2117 -- > maxViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((5,"a"), singleton 3 "b")
2118 -- > maxViewWithKey empty == Nothing
2119
2120 maxViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
2121 maxViewWithKey t = case t of
2122 Nil -> Nothing
2123 _ -> Just $ case maxViewWithKeySure t of
2124 View k v t' -> ((k, v), t')
2125 {-# INLINE maxViewWithKey #-}
2126
2127 maxViewWithKeySure :: IntMap a -> View a
2128 maxViewWithKeySure t =
2129 case t of
2130 Nil -> error "maxViewWithKeySure Nil"
2131 Bin p m l r | m < 0 ->
2132 case go l of View k a l' -> View k a (binCheckLeft p m l' r)
2133 _ -> go t
2134 where
2135 go (Bin p m l r) =
2136 case go r of View k a r' -> View k a (binCheckRight p m l r')
2137 go (Tip k y) = View k y Nil
2138 go Nil = error "maxViewWithKey_go Nil"
2139 -- See note on NOINLINE at minViewWithKeySure
2140 {-# NOINLINE maxViewWithKeySure #-}
2141
2142 -- | /O(min(n,W))/. Retrieves the minimal (key,value) pair of the map, and
2143 -- the map stripped of that element, or 'Nothing' if passed an empty map.
2144 --
2145 -- > minViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((3,"b"), singleton 5 "a")
2146 -- > minViewWithKey empty == Nothing
2147
2148 minViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
2149 minViewWithKey t =
2150 case t of
2151 Nil -> Nothing
2152 _ -> Just $ case minViewWithKeySure t of
2153 View k v t' -> ((k, v), t')
2154 -- We inline this to give GHC the best possible chance of
2155 -- getting rid of the Maybe, pair, and Int constructors, as
2156 -- well as a thunk under the Just. That is, we really want to
2157 -- be certain this inlines!
2158 {-# INLINE minViewWithKey #-}
2159
2160 minViewWithKeySure :: IntMap a -> View a
2161 minViewWithKeySure t =
2162 case t of
2163 Nil -> error "minViewWithKeySure Nil"
2164 Bin p m l r | m < 0 ->
2165 case go r of
2166 View k a r' -> View k a (binCheckRight p m l r')
2167 _ -> go t
2168 where
2169 go (Bin p m l r) =
2170 case go l of View k a l' -> View k a (binCheckLeft p m l' r)
2171 go (Tip k y) = View k y Nil
2172 go Nil = error "minViewWithKey_go Nil"
2173 -- There's never anything significant to be gained by inlining
2174 -- this. Sufficiently recent GHC versions will inline the wrapper
2175 -- anyway, which should be good enough.
2176 {-# NOINLINE minViewWithKeySure #-}
2177
2178 -- | /O(min(n,W))/. Update the value at the maximal key.
2179 --
2180 -- > updateMax (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "Xa")]
2181 -- > updateMax (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
2182
2183 updateMax :: (a -> Maybe a) -> IntMap a -> IntMap a
2184 updateMax f = updateMaxWithKey (const f)
2185
2186 -- | /O(min(n,W))/. Update the value at the minimal key.
2187 --
2188 -- > updateMin (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "Xb"), (5, "a")]
2189 -- > updateMin (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
2190
2191 updateMin :: (a -> Maybe a) -> IntMap a -> IntMap a
2192 updateMin f = updateMinWithKey (const f)
2193
2194 -- | /O(min(n,W))/. Retrieves the maximal key of the map, and the map
2195 -- stripped of that element, or 'Nothing' if passed an empty map.
2196 maxView :: IntMap a -> Maybe (a, IntMap a)
2197 maxView t = fmap (\((_, x), t') -> (x, t')) (maxViewWithKey t)
2198
2199 -- | /O(min(n,W))/. Retrieves the minimal key of the map, and the map
2200 -- stripped of that element, or 'Nothing' if passed an empty map.
2201 minView :: IntMap a -> Maybe (a, IntMap a)
2202 minView t = fmap (\((_, x), t') -> (x, t')) (minViewWithKey t)
2203
2204 -- | /O(min(n,W))/. Delete and find the maximal element.
2205 -- This function throws an error if the map is empty. Use 'maxViewWithKey'
2206 -- if the map may be empty.
2207 deleteFindMax :: IntMap a -> ((Key, a), IntMap a)
2208 deleteFindMax = fromMaybe (error "deleteFindMax: empty map has no maximal element") . maxViewWithKey
2209
2210 -- | /O(min(n,W))/. Delete and find the minimal element.
2211 -- This function throws an error if the map is empty. Use 'minViewWithKey'
2212 -- if the map may be empty.
2213 deleteFindMin :: IntMap a -> ((Key, a), IntMap a)
2214 deleteFindMin = fromMaybe (error "deleteFindMin: empty map has no minimal element") . minViewWithKey
2215
2216 -- | /O(min(n,W))/. The minimal key of the map. Returns 'Nothing' if the map is empty.
2217 lookupMin :: IntMap a -> Maybe (Key, a)
2218 lookupMin Nil = Nothing
2219 lookupMin (Tip k v) = Just (k,v)
2220 lookupMin (Bin _ m l r)
2221 | m < 0 = go r
2222 | otherwise = go l
2223 where go (Tip k v) = Just (k,v)
2224 go (Bin _ _ l' _) = go l'
2225 go Nil = Nothing
2226
2227 -- | /O(min(n,W))/. The minimal key of the map. Calls 'error' if the map is empty.
2228 -- Use 'minViewWithKey' if the map may be empty.
2229 findMin :: IntMap a -> (Key, a)
2230 findMin t
2231 | Just r <- lookupMin t = r
2232 | otherwise = error "findMin: empty map has no minimal element"
2233
2234 -- | /O(min(n,W))/. The maximal key of the map. Returns 'Nothing' if the map is empty.
2235 lookupMax :: IntMap a -> Maybe (Key, a)
2236 lookupMax Nil = Nothing
2237 lookupMax (Tip k v) = Just (k,v)
2238 lookupMax (Bin _ m l r)
2239 | m < 0 = go l
2240 | otherwise = go r
2241 where go (Tip k v) = Just (k,v)
2242 go (Bin _ _ _ r') = go r'
2243 go Nil = Nothing
2244
2245 -- | /O(min(n,W))/. The maximal key of the map. Calls 'error' if the map is empty.
2246 -- Use 'maxViewWithKey' if the map may be empty.
2247 findMax :: IntMap a -> (Key, a)
2248 findMax t
2249 | Just r <- lookupMax t = r
2250 | otherwise = error "findMax: empty map has no maximal element"
2251
2252 -- | /O(min(n,W))/. Delete the minimal key. Returns an empty map if the map is empty.
2253 --
2254 -- Note that this is a change of behaviour for consistency with 'Data.Map.Map' &#8211;
2255 -- versions prior to 0.5 threw an error if the 'IntMap' was already empty.
2256 deleteMin :: IntMap a -> IntMap a
2257 deleteMin = maybe Nil snd . minView
2258
2259 -- | /O(min(n,W))/. Delete the maximal key. Returns an empty map if the map is empty.
2260 --
2261 -- Note that this is a change of behaviour for consistency with 'Data.Map.Map' &#8211;
2262 -- versions prior to 0.5 threw an error if the 'IntMap' was already empty.
2263 deleteMax :: IntMap a -> IntMap a
2264 deleteMax = maybe Nil snd . maxView
2265
2266
2267 {--------------------------------------------------------------------
2268 Submap
2269 --------------------------------------------------------------------}
2270 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
2271 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
2272 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
2273 isProperSubmapOf m1 m2
2274 = isProperSubmapOfBy (==) m1 m2
2275
2276 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
2277 The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
2278 @m1@ and @m2@ are not equal,
2279 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
2280 applied to their respective values. For example, the following
2281 expressions are all 'True':
2282
2283 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
2284 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
2285
2286 But the following are all 'False':
2287
2288 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
2289 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
2290 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
2291 -}
2292 isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
2293 isProperSubmapOfBy predicate t1 t2
2294 = case submapCmp predicate t1 t2 of
2295 LT -> True
2296 _ -> False
2297
2298 submapCmp :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Ordering
2299 submapCmp predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
2300 | shorter m1 m2 = GT
2301 | shorter m2 m1 = submapCmpLt
2302 | p1 == p2 = submapCmpEq
2303 | otherwise = GT -- disjoint
2304 where
2305 submapCmpLt | nomatch p1 p2 m2 = GT
2306 | zero p1 m2 = submapCmp predicate t1 l2
2307 | otherwise = submapCmp predicate t1 r2
2308 submapCmpEq = case (submapCmp predicate l1 l2, submapCmp predicate r1 r2) of
2309 (GT,_ ) -> GT
2310 (_ ,GT) -> GT
2311 (EQ,EQ) -> EQ
2312 _ -> LT
2313
2314 submapCmp _ (Bin _ _ _ _) _ = GT
2315 submapCmp predicate (Tip kx x) (Tip ky y)
2316 | (kx == ky) && predicate x y = EQ
2317 | otherwise = GT -- disjoint
2318 submapCmp predicate (Tip k x) t
2319 = case lookup k t of
2320 Just y | predicate x y -> LT
2321 _ -> GT -- disjoint
2322 submapCmp _ Nil Nil = EQ
2323 submapCmp _ Nil _ = LT
2324
2325 -- | /O(n+m)/. Is this a submap?
2326 -- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
2327 isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
2328 isSubmapOf m1 m2
2329 = isSubmapOfBy (==) m1 m2
2330
2331 {- | /O(n+m)/.
2332 The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if
2333 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
2334 applied to their respective values. For example, the following
2335 expressions are all 'True':
2336
2337 > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
2338 > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
2339 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
2340
2341 But the following are all 'False':
2342
2343 > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)])
2344 > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
2345 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
2346 -}
2347 isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
2348 isSubmapOfBy predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
2349 | shorter m1 m2 = False
2350 | shorter m2 m1 = match p1 p2 m2 &&
2351 if zero p1 m2
2352 then isSubmapOfBy predicate t1 l2
2353 else isSubmapOfBy predicate t1 r2
2354 | otherwise = (p1==p2) && isSubmapOfBy predicate l1 l2 && isSubmapOfBy predicate r1 r2
2355 isSubmapOfBy _ (Bin _ _ _ _) _ = False
2356 isSubmapOfBy predicate (Tip k x) t = case lookup k t of
2357 Just y -> predicate x y
2358 Nothing -> False
2359 isSubmapOfBy _ Nil _ = True
2360
2361 {--------------------------------------------------------------------
2362 Mapping
2363 --------------------------------------------------------------------}
2364 -- | /O(n)/. Map a function over all values in the map.
2365 --
2366 -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")]
2367
2368 map :: (a -> b) -> IntMap a -> IntMap b
2369 map f = go
2370 where
2371 go (Bin p m l r) = Bin p m (go l) (go r)
2372 go (Tip k x) = Tip k (f x)
2373 go Nil = Nil
2374
2375 #ifdef __GLASGOW_HASKELL__
2376 {-# NOINLINE [1] map #-}
2377 {-# RULES
2378 "map/map" forall f g xs . map f (map g xs) = map (f . g) xs
2379 #-}
2380 #endif
2381 #if __GLASGOW_HASKELL__ >= 709
2382 -- Safe coercions were introduced in 7.8, but did not play well with RULES yet.
2383 {-# RULES
2384 "map/coerce" map coerce = coerce
2385 #-}
2386 #endif
2387
2388 -- | /O(n)/. Map a function over all values in the map.
2389 --
2390 -- > let f key x = (show key) ++ ":" ++ x
2391 -- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")]
2392
2393 mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b
2394 mapWithKey f t
2395 = case t of
2396 Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r)
2397 Tip k x -> Tip k (f k x)
2398 Nil -> Nil
2399
2400 #ifdef __GLASGOW_HASKELL__
2401 {-# NOINLINE [1] mapWithKey #-}
2402 {-# RULES
2403 "mapWithKey/mapWithKey" forall f g xs . mapWithKey f (mapWithKey g xs) =
2404 mapWithKey (\k a -> f k (g k a)) xs
2405 "mapWithKey/map" forall f g xs . mapWithKey f (map g xs) =
2406 mapWithKey (\k a -> f k (g a)) xs
2407 "map/mapWithKey" forall f g xs . map f (mapWithKey g xs) =
2408 mapWithKey (\k a -> f (g k a)) xs
2409 #-}
2410 #endif
2411
2412 -- | /O(n)/.
2413 -- @'traverseWithKey' f s == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@
2414 -- That is, behaves exactly like a regular 'traverse' except that the traversing
2415 -- function also has access to the key associated with a value.
2416 --
2417 -- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(1, 'a'), (5, 'e')]) == Just (fromList [(1, 'b'), (5, 'f')])
2418 -- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(2, 'c')]) == Nothing
2419 traverseWithKey :: Applicative t => (Key -> a -> t b) -> IntMap a -> t (IntMap b)
2420 traverseWithKey f = go
2421 where
2422 go Nil = pure Nil
2423 go (Tip k v) = Tip k <$> f k v
2424 go (Bin p m l r)
2425 | m < 0 = liftA2 (Bin p m) (go r) (go l)
2426 | otherwise = liftA2 (Bin p m) (go l) (go r)
2427 {-# INLINE traverseWithKey #-}
2428
2429 -- | /O(n)/. The function @'mapAccum'@ threads an accumulating
2430 -- argument through the map in ascending order of keys.
2431 --
2432 -- > let f a b = (a ++ b, b ++ "X")
2433 -- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")])
2434
2435 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
2436 mapAccum f = mapAccumWithKey (\a' _ x -> f a' x)
2437
2438 -- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating
2439 -- argument through the map in ascending order of keys.
2440 --
2441 -- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X")
2442 -- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")])
2443
2444 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
2445 mapAccumWithKey f a t
2446 = mapAccumL f a t
2447
2448 -- | /O(n)/. The function @'mapAccumL'@ threads an accumulating
2449 -- argument through the map in ascending order of keys.
2450 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
2451 mapAccumL f a t
2452 = case t of
2453 Bin p m l r -> let (a1,l') = mapAccumL f a l
2454 (a2,r') = mapAccumL f a1 r
2455 in (a2,Bin p m l' r')
2456 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
2457 Nil -> (a,Nil)
2458
2459 -- | /O(n)/. The function @'mapAccumR'@ threads an accumulating
2460 -- argument through the map in descending order of keys.
2461 mapAccumRWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
2462 mapAccumRWithKey f a t
2463 = case t of
2464 Bin p m l r -> let (a1,r') = mapAccumRWithKey f a r
2465 (a2,l') = mapAccumRWithKey f a1 l
2466 in (a2,Bin p m l' r')
2467 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
2468 Nil -> (a,Nil)
2469
2470 -- | /O(n*min(n,W))/.
2471 -- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@.
2472 --
2473 -- The size of the result may be smaller if @f@ maps two or more distinct
2474 -- keys to the same new key. In this case the value at the greatest of the
2475 -- original keys is retained.
2476 --
2477 -- > mapKeys (+ 1) (fromList [(5,"a"), (3,"b")]) == fromList [(4, "b"), (6, "a")]
2478 -- > mapKeys (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "c"
2479 -- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c"
2480
2481 mapKeys :: (Key->Key) -> IntMap a -> IntMap a
2482 mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) []
2483
2484 -- | /O(n*min(n,W))/.
2485 -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@.
2486 --
2487 -- The size of the result may be smaller if @f@ maps two or more distinct
2488 -- keys to the same new key. In this case the associated values will be
2489 -- combined using @c@.
2490 --
2491 -- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab"
2492 -- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab"
2493
2494 mapKeysWith :: (a -> a -> a) -> (Key->Key) -> IntMap a -> IntMap a
2495 mapKeysWith c f
2496 = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) []
2497
2498 -- | /O(n*min(n,W))/.
2499 -- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@
2500 -- is strictly monotonic.
2501 -- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@.
2502 -- /The precondition is not checked./
2503 -- Semi-formally, we have:
2504 --
2505 -- > and [x < y ==> f x < f y | x <- ls, y <- ls]
2506 -- > ==> mapKeysMonotonic f s == mapKeys f s
2507 -- > where ls = keys s
2508 --
2509 -- This means that @f@ maps distinct original keys to distinct resulting keys.
2510 -- This function has slightly better performance than 'mapKeys'.
2511 --
2512 -- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")]
2513
2514 mapKeysMonotonic :: (Key->Key) -> IntMap a -> IntMap a
2515 mapKeysMonotonic f
2516 = fromDistinctAscList . foldrWithKey (\k x xs -> (f k, x) : xs) []
2517
2518 {--------------------------------------------------------------------
2519 Filter
2520 --------------------------------------------------------------------}
2521 -- | /O(n)/. Filter all values that satisfy some predicate.
2522 --
2523 -- > filter (> "a") (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
2524 -- > filter (> "x") (fromList [(5,"a"), (3,"b")]) == empty
2525 -- > filter (< "a") (fromList [(5,"a"), (3,"b")]) == empty
2526
2527 filter :: (a -> Bool) -> IntMap a -> IntMap a
2528 filter p m
2529 = filterWithKey (\_ x -> p x) m
2530
2531 -- | /O(n)/. Filter all keys\/values that satisfy some predicate.
2532 --
2533 -- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
2534
2535 filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a
2536 filterWithKey predicate = go
2537 where
2538 go Nil = Nil
2539 go t@(Tip k x) = if predicate k x then t else Nil
2540 go (Bin p m l r) = bin p m (go l) (go r)
2541
2542 -- | /O(n)/. Partition the map according to some predicate. The first
2543 -- map contains all elements that satisfy the predicate, the second all
2544 -- elements that fail the predicate. See also 'split'.
2545 --
2546 -- > partition (> "a") (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
2547 -- > partition (< "x") (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
2548 -- > partition (> "x") (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
2549
2550 partition :: (a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
2551 partition p m
2552 = partitionWithKey (\_ x -> p x) m
2553
2554 -- | /O(n)/. Partition the map according to some predicate. The first
2555 -- map contains all elements that satisfy the predicate, the second all
2556 -- elements that fail the predicate. See also 'split'.
2557 --
2558 -- > partitionWithKey (\ k _ -> k > 3) (fromList [(5,"a"), (3,"b")]) == (singleton 5 "a", singleton 3 "b")
2559 -- > partitionWithKey (\ k _ -> k < 7) (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
2560 -- > partitionWithKey (\ k _ -> k > 7) (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
2561
2562 partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
2563 partitionWithKey predicate0 t0 = toPair $ go predicate0 t0
2564 where
2565 go predicate t =
2566 case t of
2567 Bin p m l r ->
2568 let (l1 :*: l2) = go predicate l
2569 (r1 :*: r2) = go predicate r
2570 in bin p m l1 r1 :*: bin p m l2 r2
2571 Tip k x
2572 | predicate k x -> (t :*: Nil)
2573 | otherwise -> (Nil :*: t)
2574 Nil -> (Nil :*: Nil)
2575
2576 -- | /O(n)/. Map values and collect the 'Just' results.
2577 --
2578 -- > let f x = if x == "a" then Just "new a" else Nothing
2579 -- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a"
2580
2581 mapMaybe :: (a -> Maybe b) -> IntMap a -> IntMap b
2582 mapMaybe f = mapMaybeWithKey (\_ x -> f x)
2583
2584 -- | /O(n)/. Map keys\/values and collect the 'Just' results.
2585 --
2586 -- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing
2587 -- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3"
2588
2589 mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b
2590 mapMaybeWithKey f (Bin p m l r)
2591 = bin p m (mapMaybeWithKey f l) (mapMaybeWithKey f r)
2592 mapMaybeWithKey f (Tip k x) = case f k x of
2593 Just y -> Tip k y
2594 Nothing -> Nil
2595 mapMaybeWithKey _ Nil = Nil
2596
2597 -- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
2598 --
2599 -- > let f a = if a < "c" then Left a else Right a
2600 -- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
2601 -- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")])
2602 -- >
2603 -- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
2604 -- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
2605
2606 mapEither :: (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
2607 mapEither f m
2608 = mapEitherWithKey (\_ x -> f x) m
2609
2610 -- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
2611 --
2612 -- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a)
2613 -- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
2614 -- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")])
2615 -- >
2616 -- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
2617 -- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")])
2618
2619 mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
2620 mapEitherWithKey f0 t0 = toPair $ go f0 t0
2621 where
2622 go f (Bin p m l r) =
2623 bin p m l1 r1 :*: bin p m l2 r2
2624 where
2625 (l1 :*: l2) = go f l
2626 (r1 :*: r2) = go f r
2627 go f (Tip k x) = case f k x of
2628 Left y -> (Tip k y :*: Nil)
2629 Right z -> (Nil :*: Tip k z)
2630 go _ Nil = (Nil :*: Nil)
2631
2632 -- | /O(min(n,W))/. The expression (@'split' k map@) is a pair @(map1,map2)@
2633 -- where all keys in @map1@ are lower than @k@ and all keys in
2634 -- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
2635 --
2636 -- > split 2 (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3,"b"), (5,"a")])
2637 -- > split 3 (fromList [(5,"a"), (3,"b")]) == (empty, singleton 5 "a")
2638 -- > split 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
2639 -- > split 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", empty)
2640 -- > split 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], empty)
2641
2642 split :: Key -> IntMap a -> (IntMap a, IntMap a)
2643 split k t =
2644 case t of
2645 Bin _ m l r
2646 | m < 0 ->
2647 if k >= 0 -- handle negative numbers.
2648 then
2649 case go k l of
2650 (lt :*: gt) ->
2651 let !lt' = union r lt
2652 in (lt', gt)
2653 else
2654 case go k r of
2655 (lt :*: gt) ->
2656 let !gt' = union gt l
2657 in (lt, gt')
2658 _ -> case go k t of
2659 (lt :*: gt) -> (lt, gt)
2660 where
2661 go k' t'@(Bin p m l r)
2662 | nomatch k' p m = if k' > p then t' :*: Nil else Nil :*: t'
2663 | zero k' m = case go k' l of (lt :*: gt) -> lt :*: union gt r
2664 | otherwise = case go k' r of (lt :*: gt) -> union l lt :*: gt
2665 go k' t'@(Tip ky _)
2666 | k' > ky = (t' :*: Nil)
2667 | k' < ky = (Nil :*: t')
2668 | otherwise = (Nil :*: Nil)
2669 go _ Nil = (Nil :*: Nil)
2670
2671
2672 data SplitLookup a = SplitLookup !(IntMap a) !(Maybe a) !(IntMap a)
2673
2674 mapLT :: (IntMap a -> IntMap a) -> SplitLookup a -> SplitLookup a
2675 mapLT f (SplitLookup lt fnd gt) = SplitLookup (f lt) fnd gt
2676 {-# INLINE mapLT #-}
2677
2678 mapGT :: (IntMap a -> IntMap a) -> SplitLookup a -> SplitLookup a
2679 mapGT f (SplitLookup lt fnd gt) = SplitLookup lt fnd (f gt)
2680 {-# INLINE mapGT #-}
2681
2682 -- | /O(min(n,W))/. Performs a 'split' but also returns whether the pivot
2683 -- key was found in the original map.
2684 --
2685 -- > splitLookup 2 (fromList [(5,"a"), (3,"b")]) == (empty, Nothing, fromList [(3,"b"), (5,"a")])
2686 -- > splitLookup 3 (fromList [(5,"a"), (3,"b")]) == (empty, Just "b", singleton 5 "a")
2687 -- > splitLookup 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Nothing, singleton 5 "a")
2688 -- > splitLookup 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Just "a", empty)
2689 -- > splitLookup 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], Nothing, empty)
2690
2691 splitLookup :: Key -> IntMap a -> (IntMap a, Maybe a, IntMap a)
2692 splitLookup k t =
2693 case
2694 case t of
2695 Bin _ m l r
2696 | m < 0 ->
2697 if k >= 0 -- handle negative numbers.
2698 then mapLT (union r) (go k l)
2699 else mapGT (`union` l) (go k r)
2700 _ -> go k t
2701 of SplitLookup lt fnd gt -> (lt, fnd, gt)
2702 where
2703 go k' t'@(Bin p m l r)
2704 | nomatch k' p m =
2705 if k' > p
2706 then SplitLookup t' Nothing Nil
2707 else SplitLookup Nil Nothing t'
2708 | zero k' m = mapGT (`union` r) (go k' l)
2709 | otherwise = mapLT (union l) (go k' r)
2710 go k' t'@(Tip ky y)
2711 | k' > ky = SplitLookup t' Nothing Nil
2712 | k' < ky = SplitLookup Nil Nothing t'
2713 | otherwise = SplitLookup Nil (Just y) Nil
2714 go _ Nil = SplitLookup Nil Nothing Nil
2715
2716 {--------------------------------------------------------------------
2717 Fold
2718 --------------------------------------------------------------------}
2719 -- | /O(n)/. Fold the values in the map using the given right-associative
2720 -- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'elems'@.
2721 --
2722 -- For example,
2723 --
2724 -- > elems map = foldr (:) [] map
2725 --
2726 -- > let f a len = len + (length a)
2727 -- > foldr f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
2728 foldr :: (a -> b -> b) -> b -> IntMap a -> b
2729 foldr f z = \t -> -- Use lambda t to be inlinable with two arguments only.
2730 case t of
2731 Bin _ m l r
2732 | m < 0 -> go (go z l) r -- put negative numbers before
2733 | otherwise -> go (go z r) l
2734 _ -> go z t
2735 where
2736 go z' Nil = z'
2737 go z' (Tip _ x) = f x z'
2738 go z' (Bin _ _ l r) = go (go z' r) l
2739 {-# INLINE foldr #-}
2740
2741 -- | /O(n)/. A strict version of 'foldr'. Each application of the operator is
2742 -- evaluated before using the result in the next application. This
2743 -- function is strict in the starting value.
2744 foldr' :: (a -> b -> b) -> b -> IntMap a -> b
2745 foldr' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
2746 case t of
2747 Bin _ m l r
2748 | m < 0 -> go (go z l) r -- put negative numbers before
2749 | otherwise -> go (go z r) l
2750 _ -> go z t
2751 where
2752 go !z' Nil = z'
2753 go z' (Tip _ x) = f x z'
2754 go z' (Bin _ _ l r) = go (go z' r) l
2755 {-# INLINE foldr' #-}
2756
2757 -- | /O(n)/. Fold the values in the map using the given left-associative
2758 -- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@.
2759 --
2760 -- For example,
2761 --
2762 -- > elems = reverse . foldl (flip (:)) []
2763 --
2764 -- > let f len a = len + (length a)
2765 -- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
2766 foldl :: (a -> b -> a) -> a -> IntMap b -> a
2767 foldl f z = \t -> -- Use lambda t to be inlinable with two arguments only.
2768 case t of
2769 Bin _ m l r
2770 | m < 0 -> go (go z r) l -- put negative numbers before
2771 | otherwise -> go (go z l) r
2772 _ -> go z t
2773 where
2774 go z' Nil = z'
2775 go z' (Tip _ x) = f z' x
2776 go z' (Bin _ _ l r) = go (go z' l) r
2777 {-# INLINE foldl #-}
2778
2779 -- | /O(n)/. A strict version of 'foldl'. Each application of the operator is
2780 -- evaluated before using the result in the next application. This
2781 -- function is strict in the starting value.
2782 foldl' :: (a -> b -> a) -> a -> IntMap b -> a
2783 foldl' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
2784 case t of
2785 Bin _ m l r
2786 | m < 0 -> go (go z r) l -- put negative numbers before
2787 | otherwise -> go (go z l) r
2788 _ -> go z t
2789 where
2790 go !z' Nil = z'
2791 go z' (Tip _ x) = f z' x
2792 go z' (Bin _ _ l r) = go (go z' l) r
2793 {-# INLINE foldl' #-}
2794
2795 -- | /O(n)/. Fold the keys and values in the map using the given right-associative
2796 -- binary operator, such that
2797 -- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
2798 --
2799 -- For example,
2800 --
2801 -- > keys map = foldrWithKey (\k x ks -> k:ks) [] map
2802 --
2803 -- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
2804 -- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)"
2805 foldrWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
2806 foldrWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only.
2807 case t of
2808 Bin _ m l r
2809 | m < 0 -> go (go z l) r -- put negative numbers before
2810 | otherwise -> go (go z r) l
2811 _ -> go z t
2812 where
2813 go z' Nil = z'
2814 go z' (Tip kx x) = f kx x z'
2815 go z' (Bin _ _ l r) = go (go z' r) l
2816 {-# INLINE foldrWithKey #-}
2817
2818 -- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is
2819 -- evaluated before using the result in the next application. This
2820 -- function is strict in the starting value.
2821 foldrWithKey' :: (Key -> a -> b -> b) -> b -> IntMap a -> b
2822 foldrWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
2823 case t of
2824 Bin _ m l r
2825 | m < 0 -> go (go z l) r -- put negative numbers before
2826 | otherwise -> go (go z r) l
2827 _ -> go z t
2828 where
2829 go !z' Nil = z'
2830 go z' (Tip kx x) = f kx x z'
2831 go z' (Bin _ _ l r) = go (go z' r) l
2832 {-# INLINE foldrWithKey' #-}
2833
2834 -- | /O(n)/. Fold the keys and values in the map using the given left-associative
2835 -- binary operator, such that
2836 -- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@.
2837 --
2838 -- For example,
2839 --
2840 -- > keys = reverse . foldlWithKey (\ks k x -> k:ks) []
2841 --
2842 -- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
2843 -- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)"
2844 foldlWithKey :: (a -> Key -> b -> a) -> a -> IntMap b -> a
2845 foldlWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only.
2846 case t of
2847 Bin _ m l r
2848 | m < 0 -> go (go z r) l -- put negative numbers before
2849 | otherwise -> go (go z l) r
2850 _ -> go z t
2851 where
2852 go z' Nil = z'
2853 go z' (Tip kx x) = f z' kx x
2854 go z' (Bin _ _ l r) = go (go z' l) r
2855 {-# INLINE foldlWithKey #-}
2856
2857 -- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is
2858 -- evaluated before using the result in the next application. This
2859 -- function is strict in the starting value.
2860 foldlWithKey' :: (a -> Key -> b -> a) -> a -> IntMap b -> a
2861 foldlWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
2862 case t of
2863 Bin _ m l r
2864 | m < 0 -> go (go z r) l -- put negative numbers before
2865 | otherwise -> go (go z l) r
2866 _ -> go z t
2867 where
2868 go !z' Nil = z'
2869 go z' (Tip kx x) = f z' kx x
2870 go z' (Bin _ _ l r) = go (go z' l) r
2871 {-# INLINE foldlWithKey' #-}
2872
2873 -- | /O(n)/. Fold the keys and values in the map using the given monoid, such that
2874 --
2875 -- @'foldMapWithKey' f = 'Prelude.fold' . 'mapWithKey' f@
2876 --
2877 -- This can be an asymptotically faster than 'foldrWithKey' or 'foldlWithKey' for some monoids.
2878 --
2879 -- @since 0.5.4
2880 foldMapWithKey :: Monoid m => (Key -> a -> m) -> IntMap a -> m
2881 foldMapWithKey f = go
2882 where
2883 go Nil = mempty
2884 go (Tip kx x) = f kx x
2885 go (Bin _ m l r)
2886 | m < 0 = go r `mappend` go l
2887 | otherwise = go l `mappend` go r
2888 {-# INLINE foldMapWithKey #-}
2889
2890 {--------------------------------------------------------------------
2891 List variations
2892 --------------------------------------------------------------------}
2893 -- | /O(n)/.
2894 -- Return all elements of the map in the ascending order of their keys.
2895 -- Subject to list fusion.
2896 --
2897 -- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"]
2898 -- > elems empty == []
2899
2900 elems :: IntMap a -> [a]
2901 elems = foldr (:) []
2902
2903 -- | /O(n)/. Return all keys of the map in ascending order. Subject to list
2904 -- fusion.
2905 --
2906 -- > keys (fromList [(5,"a"), (3,"b")]) == [3,5]
2907 -- > keys empty == []
2908
2909 keys :: IntMap a -> [Key]
2910 keys = foldrWithKey (\k _ ks -> k : ks) []
2911
2912 -- | /O(n)/. An alias for 'toAscList'. Returns all key\/value pairs in the
2913 -- map in ascending key order. Subject to list fusion.
2914 --
2915 -- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
2916 -- > assocs empty == []
2917
2918 assocs :: IntMap a -> [(Key,a)]
2919 assocs = toAscList
2920
2921 -- | /O(n*min(n,W))/. The set of all keys of the map.
2922 --
2923 -- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.IntSet.fromList [3,5]
2924 -- > keysSet empty == Data.IntSet.empty
2925
2926 keysSet :: IntMap a -> IntSet.IntSet
2927 keysSet Nil = IntSet.Nil
2928 keysSet (Tip kx _) = IntSet.singleton kx
2929 keysSet (Bin p m l r)
2930 | m .&. IntSet.suffixBitMask == 0 = IntSet.Bin p m (keysSet l) (keysSet r)
2931 | otherwise = IntSet.Tip (p .&. IntSet.prefixBitMask) (computeBm (computeBm 0 l) r)
2932 where computeBm !acc (Bin _ _ l' r') = computeBm (computeBm acc l') r'
2933 computeBm acc (Tip kx _) = acc .|. IntSet.bitmapOf kx
2934 computeBm _ Nil = error "Data.IntSet.keysSet: Nil"
2935
2936 -- | /O(n)/. Build a map from a set of keys and a function which for each key
2937 -- computes its value.
2938 --
2939 -- > fromSet (\k -> replicate k 'a') (Data.IntSet.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")]
2940 -- > fromSet undefined Data.IntSet.empty == empty
2941
2942 fromSet :: (Key -> a) -> IntSet.IntSet -> IntMap a
2943 fromSet _ IntSet.Nil = Nil
2944 fromSet f (IntSet.Bin p m l r) = Bin p m (fromSet f l) (fromSet f r)
2945 fromSet f (IntSet.Tip kx bm) = buildTree f kx bm (IntSet.suffixBitMask + 1)
2946 where
2947 -- This is slightly complicated, as we to convert the dense
2948 -- representation of IntSet into tree representation of IntMap.
2949 --
2950 -- We are given a nonzero bit mask 'bmask' of 'bits' bits with
2951 -- prefix 'prefix'. We split bmask into halves corresponding
2952 -- to left and right subtree. If they are both nonempty, we
2953 -- create a Bin node, otherwise exactly one of them is nonempty
2954 -- and we construct the IntMap from that half.
2955 buildTree g !prefix !bmask bits = case bits of
2956 0 -> Tip prefix (g prefix)
2957 _ -> case intFromNat ((natFromInt bits) `shiftRL` 1) of
2958 bits2
2959 | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 ->
2960 buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2
2961 | (bmask `shiftRL` bits2) .&. ((1 `shiftLL` bits2) - 1) == 0 ->
2962 buildTree g prefix bmask bits2
2963 | otherwise ->
2964 Bin prefix bits2
2965 (buildTree g prefix bmask bits2)
2966 (buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2)
2967
2968 {--------------------------------------------------------------------
2969 Lists
2970 --------------------------------------------------------------------}
2971 #if __GLASGOW_HASKELL__ >= 708
2972 -- | @since 0.5.6.2
2973 instance GHCExts.IsList (IntMap a) where
2974 type Item (IntMap a) = (Key,a)
2975 fromList = fromList
2976 toList = toList
2977 #endif
2978
2979 -- | /O(n)/. Convert the map to a list of key\/value pairs. Subject to list
2980 -- fusion.
2981 --
2982 -- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
2983 -- > toList empty == []
2984
2985 toList :: IntMap a -> [(Key,a)]
2986 toList = toAscList
2987
2988 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
2989 -- keys are in ascending order. Subject to list fusion.
2990 --
2991 -- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
2992
2993 toAscList :: IntMap a -> [(Key,a)]
2994 toAscList = foldrWithKey (\k x xs -> (k,x):xs) []
2995
2996 -- | /O(n)/. Convert the map to a list of key\/value pairs where the keys
2997 -- are in descending order. Subject to list fusion.
2998 --
2999 -- > toDescList (fromList [(5,"a"), (3,"b")]) == [(5,"a"), (3,"b")]
3000
3001 toDescList :: IntMap a -> [(Key,a)]
3002 toDescList = foldlWithKey (\xs k x -> (k,x):xs) []
3003
3004 -- List fusion for the list generating functions.
3005 #if __GLASGOW_HASKELL__
3006 -- The foldrFB and foldlFB are fold{r,l}WithKey equivalents, used for list fusion.
3007 -- They are important to convert unfused methods back, see mapFB in prelude.
3008 foldrFB :: (Key -> a -> b -> b) -> b -> IntMap a -> b
3009 foldrFB = foldrWithKey
3010 {-# INLINE[0] foldrFB #-}
3011 foldlFB :: (a -> Key -> b -> a) -> a -> IntMap b -> a
3012 foldlFB = foldlWithKey
3013 {-# INLINE[0] foldlFB #-}
3014
3015 -- Inline assocs and toList, so that we need to fuse only toAscList.
3016 {-# INLINE assocs #-}
3017 {-# INLINE toList #-}
3018
3019 -- The fusion is enabled up to phase 2 included. If it does not succeed,
3020 -- convert in phase 1 the expanded elems,keys,to{Asc,Desc}List calls back to
3021 -- elems,keys,to{Asc,Desc}List. In phase 0, we inline fold{lr}FB (which were
3022 -- used in a list fusion, otherwise it would go away in phase 1), and let compiler
3023 -- do whatever it wants with elems,keys,to{Asc,Desc}List -- it was forbidden to
3024 -- inline it before phase 0, otherwise the fusion rules would not fire at all.
3025 {-# NOINLINE[0] elems #-}
3026 {-# NOINLINE[0] keys #-}
3027 {-# NOINLINE[0] toAscList #-}
3028 {-# NOINLINE[0] toDescList #-}
3029 {-# RULES "IntMap.elems" [~1] forall m . elems m = build (\c n -> foldrFB (\_ x xs -> c x xs) n m) #-}
3030 {-# RULES "IntMap.elemsBack" [1] foldrFB (\_ x xs -> x : xs) [] = elems #-}
3031 {-# RULES "IntMap.keys" [~1] forall m . keys m = build (\c n -> foldrFB (\k _ xs -> c k xs) n m) #-}
3032 {-# RULES "IntMap.keysBack" [1] foldrFB (\k _ xs -> k : xs) [] = keys #-}
3033 {-# RULES "IntMap.toAscList" [~1] forall m . toAscList m = build (\c n -> foldrFB (\k x xs -> c (k,x) xs) n m) #-}
3034 {-# RULES "IntMap.toAscListBack" [1] foldrFB (\k x xs -> (k, x) : xs) [] = toAscList #-}
3035 {-# RULES "IntMap.toDescList" [~1] forall m . toDescList m = build (\c n -> foldlFB (\xs k x -> c (k,x) xs) n m) #-}
3036 {-# RULES "IntMap.toDescListBack" [1] foldlFB (\xs k x -> (k, x) : xs) [] = toDescList #-}
3037 #endif
3038
3039
3040 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
3041 --
3042 -- > fromList [] == empty
3043 -- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
3044 -- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
3045
3046 fromList :: [(Key,a)] -> IntMap a
3047 fromList xs
3048 = Foldable.foldl' ins empty xs
3049 where
3050 ins t (k,x) = insert k x t
3051
3052 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
3053 --
3054 -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "ab"), (5, "cba")]
3055 -- > fromListWith (++) [] == empty
3056
3057 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
3058 fromListWith f xs
3059 = fromListWithKey (\_ x y -> f x y) xs
3060
3061 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
3062 --
3063 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
3064 -- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "3:a|b"), (5, "5:c|5:b|a")]
3065 -- > fromListWithKey f [] == empty
3066
3067 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
3068 fromListWithKey f xs
3069 = Foldable.foldl' ins empty xs
3070 where
3071 ins t (k,x) = insertWithKey f k x t
3072
3073 -- | /O(n)/. Build a map from a list of key\/value pairs where
3074 -- the keys are in ascending order.
3075 --
3076 -- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
3077 -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")]
3078
3079 fromAscList :: [(Key,a)] -> IntMap a
3080 fromAscList xs
3081 = fromAscListWithKey (\_ x _ -> x) xs
3082
3083 -- | /O(n)/. Build a map from a list of key\/value pairs where
3084 -- the keys are in ascending order, with a combining function on equal keys.
3085 -- /The precondition (input list is ascending) is not checked./
3086 --
3087 -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
3088
3089 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
3090 fromAscListWith f xs
3091 = fromAscListWithKey (\_ x y -> f x y) xs
3092
3093 -- | /O(n)/. Build a map from a list of key\/value pairs where
3094 -- the keys are in ascending order, with a combining function on equal keys.
3095 -- /The precondition (input list is ascending) is not checked./
3096 --
3097 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
3098 -- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "5:b|a")]
3099
3100 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
3101 fromAscListWithKey _ [] = Nil
3102 fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0)
3103 where
3104 -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
3105 combineEq z [] = [z]
3106 combineEq z@(kz,zz) (x@(kx,xx):xs)
3107 | kx==kz = let yy = f kx xx zz in combineEq (kx,yy) xs
3108 | otherwise = z:combineEq x xs
3109
3110 -- | /O(n)/. Build a map from a list of key\/value pairs where
3111 -- the keys are in ascending order and all distinct.
3112 -- /The precondition (input list is strictly ascending) is not checked./
3113 --
3114 -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
3115
3116 #if __GLASGOW_HASKELL__
3117 fromDistinctAscList :: forall a. [(Key,a)] -> IntMap a
3118 #else
3119 fromDistinctAscList :: [(Key,a)] -> IntMap a
3120 #endif
3121 fromDistinctAscList [] = Nil
3122 fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada
3123 where
3124 work (kx,vx) [] stk = finish kx (Tip kx vx) stk
3125 work (kx,vx) (z@(kz,_):zs) stk = reduce z zs (branchMask kx kz) kx (Tip kx vx) stk
3126
3127 #if __GLASGOW_HASKELL__
3128 reduce :: (Key,a) -> [(Key,a)] -> Mask -> Prefix -> IntMap a -> Stack a -> IntMap a
3129 #endif
3130 reduce z zs _ px tx Nada = work z zs (Push px tx Nada)
3131 reduce z zs m px tx stk@(Push py ty stk') =
3132 let mxy = branchMask px py
3133 pxy = mask px mxy
3134 in if shorter m mxy
3135 then reduce z zs m pxy (Bin pxy mxy ty tx) stk'
3136 else work z zs (Push px tx stk)
3137
3138 finish _ t Nada = t
3139 finish px tx (Push py ty stk) = finish p (link py ty px tx) stk
3140 where m = branchMask px py
3141 p = mask px m
3142
3143 data Stack a = Push {-# UNPACK #-} !Prefix !(IntMap a) !(Stack a) | Nada
3144
3145
3146 {--------------------------------------------------------------------
3147 Eq
3148 --------------------------------------------------------------------}
3149 instance Eq a => Eq (IntMap a) where
3150 t1 == t2 = equal t1 t2
3151 t1 /= t2 = nequal t1 t2
3152
3153 equal :: Eq a => IntMap a -> IntMap a -> Bool
3154 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
3155 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
3156 equal (Tip kx x) (Tip ky y)
3157 = (kx == ky) && (x==y)
3158 equal Nil Nil = True
3159 equal _ _ = False
3160
3161 nequal :: Eq a => IntMap a -> IntMap a -> Bool
3162 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
3163 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
3164 nequal (Tip kx x) (Tip ky y)
3165 = (kx /= ky) || (x/=y)
3166 nequal Nil Nil = False
3167 nequal _ _ = True
3168
3169 #if MIN_VERSION_base(4,9,0)
3170 -- | @since 0.5.9
3171 instance Eq1 IntMap where
3172 liftEq eq (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
3173 = (m1 == m2) && (p1 == p2) && (liftEq eq l1 l2) && (liftEq eq r1 r2)
3174 liftEq eq (Tip kx x) (Tip ky y)
3175 = (kx == ky) && (eq x y)
3176 liftEq _eq Nil Nil = True
3177 liftEq _eq _ _ = False
3178 #endif
3179
3180 {--------------------------------------------------------------------
3181 Ord
3182 --------------------------------------------------------------------}
3183
3184 instance Ord a => Ord (IntMap a) where
3185 compare m1 m2 = compare (toList m1) (toList m2)
3186
3187 #if MIN_VERSION_base(4,9,0)
3188 -- | @since 0.5.9
3189 instance Ord1 IntMap where
3190 liftCompare cmp m n =
3191 liftCompare (liftCompare cmp) (toList m) (toList n)
3192 #endif
3193
3194 {--------------------------------------------------------------------
3195 Functor
3196 --------------------------------------------------------------------}
3197
3198 instance Functor IntMap where
3199 fmap = map
3200
3201 #ifdef __GLASGOW_HASKELL__
3202 a <$ Bin p m l r = Bin p m (a <$ l) (a <$ r)
3203 a <$ Tip k _ = Tip k a
3204 _ <$ Nil = Nil
3205 #endif
3206
3207 {--------------------------------------------------------------------
3208 Show
3209 --------------------------------------------------------------------}
3210
3211 instance Show a => Show (IntMap a) where
3212 showsPrec d m = showParen (d > 10) $
3213 showString "fromList " . shows (toList m)
3214
3215 #if MIN_VERSION_base(4,9,0)
3216 -- | @since 0.5.9
3217 instance Show1 IntMap where
3218 liftShowsPrec sp sl d m =
3219 showsUnaryWith (liftShowsPrec sp' sl') "fromList" d (toList m)
3220 where
3221 sp' = liftShowsPrec sp sl
3222 sl' = liftShowList sp sl
3223 #endif
3224
3225 {--------------------------------------------------------------------
3226 Read
3227 --------------------------------------------------------------------}
3228 instance (Read e) => Read (IntMap e) where
3229 #ifdef __GLASGOW_HASKELL__
3230 readPrec = parens $ prec 10 $ do
3231 Ident "fromList" <- lexP
3232 xs <- readPrec
3233 return (fromList xs)
3234
3235 readListPrec = readListPrecDefault
3236 #else
3237 readsPrec p = readParen (p > 10) $ \ r -> do
3238 ("fromList",s) <- lex r
3239 (xs,t) <- reads s
3240 return (fromList xs,t)
3241 #endif
3242
3243 #if MIN_VERSION_base(4,9,0)
3244 -- | @since 0.5.9
3245 instance Read1 IntMap where
3246 liftReadsPrec rp rl = readsData $
3247 readsUnaryWith (liftReadsPrec rp' rl') "fromList" fromList
3248 where
3249 rp' = liftReadsPrec rp rl
3250 rl' = liftReadList rp rl
3251 #endif
3252
3253 {--------------------------------------------------------------------
3254 Typeable
3255 --------------------------------------------------------------------}
3256
3257 INSTANCE_TYPEABLE1(IntMap)
3258
3259 {--------------------------------------------------------------------
3260 Helpers
3261 --------------------------------------------------------------------}
3262 {--------------------------------------------------------------------
3263 Link
3264 --------------------------------------------------------------------}
3265 link :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
3266 link p1 t1 p2 t2
3267 | zero p1 m = Bin p m t1 t2
3268 | otherwise = Bin p m t2 t1
3269 where
3270 m = branchMask p1 p2
3271 p = mask p1 m
3272 {-# INLINE link #-}
3273
3274 {--------------------------------------------------------------------
3275 @bin@ assures that we never have empty trees within a tree.
3276 --------------------------------------------------------------------}
3277 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
3278 bin _ _ l Nil = l
3279 bin _ _ Nil r = r
3280 bin p m l r = Bin p m l r
3281 {-# INLINE bin #-}
3282
3283 -- binCheckLeft only checks that the left subtree is non-empty
3284 binCheckLeft :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
3285 binCheckLeft _ _ Nil r = r
3286 binCheckLeft p m l r = Bin p m l r
3287 {-# INLINE binCheckLeft #-}
3288
3289 -- binCheckRight only checks that the right subtree is non-empty
3290 binCheckRight :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
3291 binCheckRight _ _ l Nil = l
3292 binCheckRight p m l r = Bin p m l r
3293 {-# INLINE binCheckRight #-}
3294
3295 {--------------------------------------------------------------------
3296 Endian independent bit twiddling
3297 --------------------------------------------------------------------}
3298
3299 -- | Should this key follow the left subtree of a 'Bin' with switching
3300 -- bit @m@? N.B., the answer is only valid when @match i p m@ is true.
3301 zero :: Key -> Mask -> Bool
3302 zero i m
3303 = (natFromInt i) .&. (natFromInt m) == 0
3304 {-# INLINE zero #-}
3305
3306 nomatch,match :: Key -> Prefix -> Mask -> Bool
3307
3308 -- | Does the key @i@ differ from the prefix @p@ before getting to
3309 -- the switching bit @m@?
3310 nomatch i p m
3311 = (mask i m) /= p
3312 {-# INLINE nomatch #-}
3313
3314 -- | Does the key @i@ match the prefix @p@ (up to but not including
3315 -- bit @m@)?
3316 match i p m
3317 = (mask i m) == p
3318 {-# INLINE match #-}
3319
3320
3321 -- | The prefix of key @i@ up to (but not including) the switching
3322 -- bit @m@.
3323 mask :: Key -> Mask -> Prefix
3324 mask i m
3325 = maskW (natFromInt i) (natFromInt m)
3326 {-# INLINE mask #-}
3327
3328
3329 {--------------------------------------------------------------------
3330 Big endian operations
3331 --------------------------------------------------------------------}
3332
3333 -- | The prefix of key @i@ up to (but not including) the switching
3334 -- bit @m@.
3335 maskW :: Nat -> Nat -> Prefix
3336 maskW i m
3337 = intFromNat (i .&. (complement (m-1) `xor` m))
3338 {-# INLINE maskW #-}
3339
3340 -- | Does the left switching bit specify a shorter prefix?
3341 shorter :: Mask -> Mask -> Bool
3342 shorter m1 m2
3343 = (natFromInt m1) > (natFromInt m2)
3344 {-# INLINE shorter #-}
3345
3346 -- | The first switching bit where the two prefixes disagree.
3347 branchMask :: Prefix -> Prefix -> Mask
3348 branchMask p1 p2
3349 = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
3350 {-# INLINE branchMask #-}
3351
3352 {--------------------------------------------------------------------
3353 Utilities
3354 --------------------------------------------------------------------}
3355
3356 -- | /O(1)/. Decompose a map into pieces based on the structure
3357 -- of the underlying tree. This function is useful for consuming a
3358 -- map in parallel.
3359 --
3360 -- No guarantee is made as to the sizes of the pieces; an internal, but
3361 -- deterministic process determines this. However, it is guaranteed that the
3362 -- pieces returned will be in ascending order (all elements in the first submap
3363 -- less than all elements in the second, and so on).
3364 --
3365 -- Examples:
3366 --
3367 -- > splitRoot (fromList (zip [1..6::Int] ['a'..])) ==
3368 -- > [fromList [(1,'a'),(2,'b'),(3,'c')],fromList [(4,'d'),(5,'e'),(6,'f')]]
3369 --
3370 -- > splitRoot empty == []
3371 --
3372 -- Note that the current implementation does not return more than two submaps,
3373 -- but you should not depend on this behaviour because it can change in the
3374 -- future without notice.
3375 splitRoot :: IntMap a -> [IntMap a]
3376 splitRoot orig =
3377 case orig of
3378 Nil -> []
3379 x@(Tip _ _) -> [x]
3380 Bin _ m l r | m < 0 -> [r, l]
3381 | otherwise -> [l, r]
3382 {-# INLINE splitRoot #-}
3383
3384
3385 {--------------------------------------------------------------------
3386 Debugging
3387 --------------------------------------------------------------------}
3388
3389 -- | /O(n)/. Show the tree that implements the map. The tree is shown
3390 -- in a compressed, hanging format.
3391 showTree :: Show a => IntMap a -> String
3392 showTree s
3393 = showTreeWith True False s
3394
3395
3396 {- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows
3397 the tree that implements the map. If @hang@ is
3398 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
3399 @wide@ is 'True', an extra wide version is shown.
3400 -}
3401 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
3402 showTreeWith hang wide t
3403 | hang = (showsTreeHang wide [] t) ""
3404 | otherwise = (showsTree wide [] [] t) ""
3405
3406 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
3407 showsTree wide lbars rbars t = case t of
3408 Bin p m l r ->
3409 showsTree wide (withBar rbars) (withEmpty rbars) r .
3410 showWide wide rbars .
3411 showsBars lbars . showString (showBin p m) . showString "\n" .
3412 showWide wide lbars .
3413 showsTree wide (withEmpty lbars) (withBar lbars) l
3414 Tip k x ->
3415 showsBars lbars .
3416 showString " " . shows k . showString ":=" . shows x . showString "\n"
3417 Nil -> showsBars lbars . showString "|\n"
3418
3419 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
3420 showsTreeHang wide bars t = case t of
3421 Bin p m l r ->
3422 showsBars bars . showString (showBin p m) . showString "\n" .
3423 showWide wide bars .
3424 showsTreeHang wide (withBar bars) l .
3425 showWide wide bars .
3426 showsTreeHang wide (withEmpty bars) r
3427 Tip k x ->
3428 showsBars bars .
3429 showString " " . shows k . showString ":=" . shows x . showString "\n"
3430 Nil -> showsBars bars . showString "|\n"
3431
3432 showBin :: Prefix -> Mask -> String
3433 showBin _ _
3434 = "*" -- ++ show (p,m)
3435
3436 showWide :: Bool -> [String] -> String -> String
3437 showWide wide bars
3438 | wide = showString (concat (reverse bars)) . showString "|\n"
3439 | otherwise = id
3440
3441 showsBars :: [String] -> ShowS
3442 showsBars bars
3443 = case bars of
3444 [] -> id
3445 _ -> showString (concat (reverse (tail bars))) . showString node
3446
3447 node :: String
3448 node = "+--"
3449
3450 withBar, withEmpty :: [String] -> [String]
3451 withBar bars = "| ":bars
3452 withEmpty bars = " ":bars