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