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