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