f8b18a3165ed14e291e3dea22eabdb33c7bbb1d3
[packages/containers.git] / Data / Map / Internal.hs
1 {-# LANGUAGE CPP #-}
2 {-# LANGUAGE BangPatterns #-}
3 {-# LANGUAGE PatternGuards #-}
4 #if __GLASGOW_HASKELL__
5 {-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-}
6 #endif
7 #if defined(__GLASGOW_HASKELL__)
8 {-# LANGUAGE Trustworthy #-}
9 #endif
10 #if __GLASGOW_HASKELL__ >= 708
11 {-# LANGUAGE RoleAnnotations #-}
12 {-# LANGUAGE TypeFamilies #-}
13 #define USE_MAGIC_PROXY 1
14 #endif
15
16 #ifdef USE_MAGIC_PROXY
17 {-# LANGUAGE MagicHash #-}
18 #endif
19
20 {-# OPTIONS_HADDOCK not-home #-}
21
22 #include "containers.h"
23
24 #if !(WORD_SIZE_IN_BITS >= 61)
25 #define DEFINE_ALTERF_FALLBACK 1
26 #endif
27
28 -----------------------------------------------------------------------------
29 -- |
30 -- Module : Data.Map.Internal
31 -- Copyright : (c) Daan Leijen 2002
32 -- (c) Andriy Palamarchuk 2008
33 -- License : BSD-style
34 -- Maintainer : libraries@haskell.org
35 -- Portability : portable
36 --
37 -- = WARNING
38 --
39 -- This module is considered __internal__.
40 --
41 -- The Package Versioning Policy __does not apply__.
42 --
43 -- The contents of this module may change __in any way whatsoever__
44 -- and __without any warning__ between minor versions of this package.
45 --
46 -- Authors importing this module are expected to track development
47 -- closely.
48 --
49 -- = Description
50 --
51 -- An efficient implementation of maps from keys to values (dictionaries).
52 --
53 -- Since many function names (but not the type name) clash with
54 -- "Prelude" names, this module is usually imported @qualified@, e.g.
55 --
56 -- > import Data.Map (Map)
57 -- > import qualified Data.Map as Map
58 --
59 -- The implementation of 'Map' is based on /size balanced/ binary trees (or
60 -- trees of /bounded balance/) as described by:
61 --
62 -- * Stephen Adams, \"/Efficient sets: a balancing act/\",
63 -- Journal of Functional Programming 3(4):553-562, October 1993,
64 -- <http://www.swiss.ai.mit.edu/~adams/BB/>.
65 -- * J. Nievergelt and E.M. Reingold,
66 -- \"/Binary search trees of bounded balance/\",
67 -- SIAM journal of computing 2(1), March 1973.
68 --
69 -- Bounds for 'union', 'intersection', and 'difference' are as given
70 -- by
71 --
72 -- * Guy Blelloch, Daniel Ferizovic, and Yihan Sun,
73 -- \"/Just Join for Parallel Ordered Sets/\",
74 -- <https://arxiv.org/abs/1602.02120v3>.
75 --
76 -- Note that the implementation is /left-biased/ -- the elements of a
77 -- first argument are always preferred to the second, for example in
78 -- 'union' or 'insert'.
79 --
80 -- Operation comments contain the operation time complexity in
81 -- the Big-O notation <http://en.wikipedia.org/wiki/Big_O_notation>.
82 --
83 -- @since 0.5.9
84 -----------------------------------------------------------------------------
85
86 -- [Note: Using INLINABLE]
87 -- ~~~~~~~~~~~~~~~~~~~~~~~
88 -- It is crucial to the performance that the functions specialize on the Ord
89 -- type when possible. GHC 7.0 and higher does this by itself when it sees th
90 -- unfolding of a function -- that is why all public functions are marked
91 -- INLINABLE (that exposes the unfolding).
92
93
94 -- [Note: Using INLINE]
95 -- ~~~~~~~~~~~~~~~~~~~~
96 -- For other compilers and GHC pre 7.0, we mark some of the functions INLINE.
97 -- We mark the functions that just navigate down the tree (lookup, insert,
98 -- delete and similar). That navigation code gets inlined and thus specialized
99 -- when possible. There is a price to pay -- code growth. The code INLINED is
100 -- therefore only the tree navigation, all the real work (rebalancing) is not
101 -- INLINED by using a NOINLINE.
102 --
103 -- All methods marked INLINE have to be nonrecursive -- a 'go' function doing
104 -- the real work is provided.
105
106
107 -- [Note: Type of local 'go' function]
108 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
109 -- If the local 'go' function uses an Ord class, it sometimes heap-allocates
110 -- the Ord dictionary when the 'go' function does not have explicit type.
111 -- In that case we give 'go' explicit type. But this slightly decrease
112 -- performance, as the resulting 'go' function can float out to top level.
113
114
115 -- [Note: Local 'go' functions and capturing]
116 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
117 -- As opposed to Map, when 'go' function captures an argument, increased
118 -- heap-allocation can occur: sometimes in a polymorphic function, the 'go'
119 -- floats out of its enclosing function and then it heap-allocates the
120 -- dictionary and the argument. Maybe it floats out too late and strictness
121 -- analyzer cannot see that these could be passed on stack.
122 --
123
124 -- [Note: Order of constructors]
125 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
126 -- The order of constructors of Map matters when considering performance.
127 -- Currently in GHC 7.0, when type has 2 constructors, a forward conditional
128 -- jump is made when successfully matching second constructor. Successful match
129 -- of first constructor results in the forward jump not taken.
130 -- On GHC 7.0, reordering constructors from Tip | Bin to Bin | Tip
131 -- improves the benchmark by up to 10% on x86.
132
133 module Data.Map.Internal (
134 -- * Map type
135 Map(..) -- instance Eq,Show,Read
136 , Size
137
138 -- * Operators
139 , (!), (!?), (\\)
140
141 -- * Query
142 , null
143 , size
144 , member
145 , notMember
146 , lookup
147 , findWithDefault
148 , lookupLT
149 , lookupGT
150 , lookupLE
151 , lookupGE
152
153 -- * Construction
154 , empty
155 , singleton
156
157 -- ** Insertion
158 , insert
159 , insertWith
160 , insertWithKey
161 , insertLookupWithKey
162
163 -- ** Delete\/Update
164 , delete
165 , adjust
166 , adjustWithKey
167 , update
168 , updateWithKey
169 , updateLookupWithKey
170 , alter
171 , alterF
172
173 -- * Combine
174
175 -- ** Union
176 , union
177 , unionWith
178 , unionWithKey
179 , unions
180 , unionsWith
181
182 -- ** Difference
183 , difference
184 , differenceWith
185 , differenceWithKey
186
187 -- ** Intersection
188 , intersection
189 , intersectionWith
190 , intersectionWithKey
191
192 -- ** General combining function
193 , SimpleWhenMissing
194 , SimpleWhenMatched
195 , runWhenMatched
196 , runWhenMissing
197 , merge
198 -- *** @WhenMatched@ tactics
199 , zipWithMaybeMatched
200 , zipWithMatched
201 -- *** @WhenMissing@ tactics
202 , mapMaybeMissing
203 , dropMissing
204 , preserveMissing
205 , preserveMissing'
206 , mapMissing
207 , filterMissing
208
209 -- ** Applicative general combining function
210 , WhenMissing (..)
211 , WhenMatched (..)
212 , mergeA
213
214 -- *** @WhenMatched@ tactics
215 -- | The tactics described for 'merge' work for
216 -- 'mergeA' as well. Furthermore, the following
217 -- are available.
218 , zipWithMaybeAMatched
219 , zipWithAMatched
220
221 -- *** @WhenMissing@ tactics
222 -- | The tactics described for 'merge' work for
223 -- 'mergeA' as well. Furthermore, the following
224 -- are available.
225 , traverseMaybeMissing
226 , traverseMissing
227 , filterAMissing
228
229 -- ** Deprecated general combining function
230
231 , mergeWithKey
232
233 -- * Traversal
234 -- ** Map
235 , map
236 , mapWithKey
237 , traverseWithKey
238 , traverseMaybeWithKey
239 , mapAccum
240 , mapAccumWithKey
241 , mapAccumRWithKey
242 , mapKeys
243 , mapKeysWith
244 , mapKeysMonotonic
245
246 -- * Folds
247 , foldr
248 , foldl
249 , foldrWithKey
250 , foldlWithKey
251 , foldMapWithKey
252
253 -- ** Strict folds
254 , foldr'
255 , foldl'
256 , foldrWithKey'
257 , foldlWithKey'
258
259 -- * Conversion
260 , elems
261 , keys
262 , assocs
263 , keysSet
264 , fromSet
265
266 -- ** Lists
267 , toList
268 , fromList
269 , fromListWith
270 , fromListWithKey
271
272 -- ** Ordered lists
273 , toAscList
274 , toDescList
275 , fromAscList
276 , fromAscListWith
277 , fromAscListWithKey
278 , fromDistinctAscList
279 , fromDescList
280 , fromDescListWith
281 , fromDescListWithKey
282 , fromDistinctDescList
283
284 -- * Filter
285 , filter
286 , filterWithKey
287
288 , takeWhileAntitone
289 , dropWhileAntitone
290 , spanAntitone
291
292 , restrictKeys
293 , withoutKeys
294 , partition
295 , partitionWithKey
296
297 , mapMaybe
298 , mapMaybeWithKey
299 , mapEither
300 , mapEitherWithKey
301
302 , split
303 , splitLookup
304 , splitRoot
305
306 -- * Submap
307 , isSubmapOf, isSubmapOfBy
308 , isProperSubmapOf, isProperSubmapOfBy
309
310 -- * Indexed
311 , lookupIndex
312 , findIndex
313 , elemAt
314 , updateAt
315 , deleteAt
316 , take
317 , drop
318 , splitAt
319
320 -- * Min\/Max
321 , lookupMin
322 , lookupMax
323 , findMin
324 , findMax
325 , deleteMin
326 , deleteMax
327 , deleteFindMin
328 , deleteFindMax
329 , updateMin
330 , updateMax
331 , updateMinWithKey
332 , updateMaxWithKey
333 , minView
334 , maxView
335 , minViewWithKey
336 , maxViewWithKey
337
338 -- Used by the strict version
339 , AreWeStrict (..)
340 , atKeyImpl
341 #if __GLASGOW_HASKELL__ && MIN_VERSION_base(4,8,0)
342 , atKeyPlain
343 #endif
344 , bin
345 , balance
346 , balanceL
347 , balanceR
348 , delta
349 , insertMax
350 , link
351 , link2
352 , glue
353 , MaybeS(..)
354 , Identity(..)
355
356 -- Used by Map.Merge.Lazy
357 , mapWhenMissing
358 , mapWhenMatched
359 , lmapWhenMissing
360 , contramapFirstWhenMatched
361 , contramapSecondWhenMatched
362 , mapGentlyWhenMissing
363 , mapGentlyWhenMatched
364 ) where
365
366 #if MIN_VERSION_base(4,8,0)
367 import Data.Functor.Identity (Identity (..))
368 import Control.Applicative (liftA3)
369 #else
370 import Control.Applicative (Applicative(..), (<$>), liftA3)
371 import Data.Monoid (Monoid(..))
372 import Data.Traversable (Traversable(traverse))
373 #endif
374 #if MIN_VERSION_base(4,9,0)
375 import Data.Functor.Classes
376 import Data.Semigroup (stimesIdempotentMonoid)
377 #endif
378 #if MIN_VERSION_base(4,9,0)
379 import Data.Semigroup (Semigroup(stimes))
380 #endif
381 #if !(MIN_VERSION_base(4,11,0)) && MIN_VERSION_base(4,9,0)
382 import Data.Semigroup (Semigroup((<>)))
383 #endif
384 import Control.Applicative (Const (..))
385 import Control.DeepSeq (NFData(rnf))
386 import Data.Bits (shiftL, shiftR)
387 import qualified Data.Foldable as Foldable
388 #if !MIN_VERSION_base(4,8,0)
389 import Data.Foldable (Foldable())
390 #endif
391 import Data.Typeable
392 import Prelude hiding (lookup, map, filter, foldr, foldl, null, splitAt, take, drop)
393
394 import qualified Data.Set.Internal as Set
395 import Data.Set.Internal (Set)
396 import Utils.Containers.Internal.PtrEquality (ptrEq)
397 import Utils.Containers.Internal.StrictPair
398 import Utils.Containers.Internal.StrictMaybe
399 import Utils.Containers.Internal.BitQueue
400 #ifdef DEFINE_ALTERF_FALLBACK
401 import Utils.Containers.Internal.BitUtil (wordSize)
402 #endif
403
404 #if __GLASGOW_HASKELL__
405 import GHC.Exts (build, lazy)
406 #if !MIN_VERSION_base(4,8,0)
407 import Data.Functor ((<$))
408 #endif
409 #ifdef USE_MAGIC_PROXY
410 import GHC.Exts (Proxy#, proxy# )
411 #endif
412 #if __GLASGOW_HASKELL__ >= 708
413 import qualified GHC.Exts as GHCExts
414 #endif
415 import Text.Read hiding (lift)
416 import Data.Data
417 import qualified Control.Category as Category
418 #endif
419 #if __GLASGOW_HASKELL__ >= 708
420 import Data.Coerce
421 #endif
422
423
424 {--------------------------------------------------------------------
425 Operators
426 --------------------------------------------------------------------}
427 infixl 9 !,!?,\\ --
428
429 -- | /O(log n)/. Find the value at a key.
430 -- Calls 'error' when the element can not be found.
431 --
432 -- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map
433 -- > fromList [(5,'a'), (3,'b')] ! 5 == 'a'
434
435 (!) :: Ord k => Map k a -> k -> a
436 (!) m k = find k m
437 #if __GLASGOW_HASKELL__
438 {-# INLINE (!) #-}
439 #endif
440
441 -- | /O(log n)/. Find the value at a key.
442 -- Returns 'Nothing' when the element can not be found.
443 --
444 -- prop> fromList [(5, 'a'), (3, 'b')] !? 1 == Nothing
445 -- prop> fromList [(5, 'a'), (3, 'b')] !? 5 == Just 'a'
446 --
447 -- @since 0.5.9
448
449 (!?) :: Ord k => Map k a -> k -> Maybe a
450 (!?) m k = lookup k m
451 #if __GLASGOW_HASKELL__
452 {-# INLINE (!?) #-}
453 #endif
454
455 -- | Same as 'difference'.
456 (\\) :: Ord k => Map k a -> Map k b -> Map k a
457 m1 \\ m2 = difference m1 m2
458 #if __GLASGOW_HASKELL__
459 {-# INLINE (\\) #-}
460 #endif
461
462 {--------------------------------------------------------------------
463 Size balanced trees.
464 --------------------------------------------------------------------}
465 -- | A Map from keys @k@ to values @a@.
466 --
467 -- The 'Semigroup' operation for 'Map' is 'union', which prefers
468 -- values from the left operand. If @m1@ maps a key @k@ to a value
469 -- @a1@, and @m2@ maps the same key to a different value @a2@, then
470 -- their union @m1 <> m2@ maps @k@ to @a1@.
471
472 -- See Note: Order of constructors
473 data Map k a = Bin {-# UNPACK #-} !Size !k a !(Map k a) !(Map k a)
474 | Tip
475
476 type Size = Int
477
478 #if __GLASGOW_HASKELL__ >= 708
479 type role Map nominal representational
480 #endif
481
482 instance (Ord k) => Monoid (Map k v) where
483 mempty = empty
484 mconcat = unions
485 #if !(MIN_VERSION_base(4,9,0))
486 mappend = union
487 #else
488 mappend = (<>)
489
490 instance (Ord k) => Semigroup (Map k v) where
491 (<>) = union
492 stimes = stimesIdempotentMonoid
493 #endif
494
495 #if __GLASGOW_HASKELL__
496
497 {--------------------------------------------------------------------
498 A Data instance
499 --------------------------------------------------------------------}
500
501 -- This instance preserves data abstraction at the cost of inefficiency.
502 -- We provide limited reflection services for the sake of data abstraction.
503
504 instance (Data k, Data a, Ord k) => Data (Map k a) where
505 gfoldl f z m = z fromList `f` toList m
506 toConstr _ = fromListConstr
507 gunfold k z c = case constrIndex c of
508 1 -> k (z fromList)
509 _ -> error "gunfold"
510 dataTypeOf _ = mapDataType
511 dataCast2 f = gcast2 f
512
513 fromListConstr :: Constr
514 fromListConstr = mkConstr mapDataType "fromList" [] Prefix
515
516 mapDataType :: DataType
517 mapDataType = mkDataType "Data.Map.Internal.Map" [fromListConstr]
518
519 #endif
520
521 {--------------------------------------------------------------------
522 Query
523 --------------------------------------------------------------------}
524 -- | /O(1)/. Is the map empty?
525 --
526 -- > Data.Map.null (empty) == True
527 -- > Data.Map.null (singleton 1 'a') == False
528
529 null :: Map k a -> Bool
530 null Tip = True
531 null (Bin {}) = False
532 {-# INLINE null #-}
533
534 -- | /O(1)/. The number of elements in the map.
535 --
536 -- > size empty == 0
537 -- > size (singleton 1 'a') == 1
538 -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3
539
540 size :: Map k a -> Int
541 size Tip = 0
542 size (Bin sz _ _ _ _) = sz
543 {-# INLINE size #-}
544
545
546 -- | /O(log n)/. Lookup the value at a key in the map.
547 --
548 -- The function will return the corresponding value as @('Just' value)@,
549 -- or 'Nothing' if the key isn't in the map.
550 --
551 -- An example of using @lookup@:
552 --
553 -- > import Prelude hiding (lookup)
554 -- > import Data.Map
555 -- >
556 -- > employeeDept = fromList([("John","Sales"), ("Bob","IT")])
557 -- > deptCountry = fromList([("IT","USA"), ("Sales","France")])
558 -- > countryCurrency = fromList([("USA", "Dollar"), ("France", "Euro")])
559 -- >
560 -- > employeeCurrency :: String -> Maybe String
561 -- > employeeCurrency name = do
562 -- > dept <- lookup name employeeDept
563 -- > country <- lookup dept deptCountry
564 -- > lookup country countryCurrency
565 -- >
566 -- > main = do
567 -- > putStrLn $ "John's currency: " ++ (show (employeeCurrency "John"))
568 -- > putStrLn $ "Pete's currency: " ++ (show (employeeCurrency "Pete"))
569 --
570 -- The output of this program:
571 --
572 -- > John's currency: Just "Euro"
573 -- > Pete's currency: Nothing
574 lookup :: Ord k => k -> Map k a -> Maybe a
575 lookup = go
576 where
577 go !_ Tip = Nothing
578 go k (Bin _ kx x l r) = case compare k kx of
579 LT -> go k l
580 GT -> go k r
581 EQ -> Just x
582 #if __GLASGOW_HASKELL__
583 {-# INLINABLE lookup #-}
584 #else
585 {-# INLINE lookup #-}
586 #endif
587
588 -- | /O(log n)/. Is the key a member of the map? See also 'notMember'.
589 --
590 -- > member 5 (fromList [(5,'a'), (3,'b')]) == True
591 -- > member 1 (fromList [(5,'a'), (3,'b')]) == False
592 member :: Ord k => k -> Map k a -> Bool
593 member = go
594 where
595 go !_ Tip = False
596 go k (Bin _ kx _ l r) = case compare k kx of
597 LT -> go k l
598 GT -> go k r
599 EQ -> True
600 #if __GLASGOW_HASKELL__
601 {-# INLINABLE member #-}
602 #else
603 {-# INLINE member #-}
604 #endif
605
606 -- | /O(log n)/. Is the key not a member of the map? See also 'member'.
607 --
608 -- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False
609 -- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True
610
611 notMember :: Ord k => k -> Map k a -> Bool
612 notMember k m = not $ member k m
613 #if __GLASGOW_HASKELL__
614 {-# INLINABLE notMember #-}
615 #else
616 {-# INLINE notMember #-}
617 #endif
618
619 -- | /O(log n)/. Find the value at a key.
620 -- Calls 'error' when the element can not be found.
621 find :: Ord k => k -> Map k a -> a
622 find = go
623 where
624 go !_ Tip = error "Map.!: given key is not an element in the map"
625 go k (Bin _ kx x l r) = case compare k kx of
626 LT -> go k l
627 GT -> go k r
628 EQ -> x
629 #if __GLASGOW_HASKELL__
630 {-# INLINABLE find #-}
631 #else
632 {-# INLINE find #-}
633 #endif
634
635 -- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns
636 -- the value at key @k@ or returns default value @def@
637 -- when the key is not in the map.
638 --
639 -- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x'
640 -- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a'
641 findWithDefault :: Ord k => a -> k -> Map k a -> a
642 findWithDefault = go
643 where
644 go def !_ Tip = def
645 go def k (Bin _ kx x l r) = case compare k kx of
646 LT -> go def k l
647 GT -> go def k r
648 EQ -> x
649 #if __GLASGOW_HASKELL__
650 {-# INLINABLE findWithDefault #-}
651 #else
652 {-# INLINE findWithDefault #-}
653 #endif
654
655 -- | /O(log n)/. Find largest key smaller than the given one and return the
656 -- corresponding (key, value) pair.
657 --
658 -- > lookupLT 3 (fromList [(3,'a'), (5,'b')]) == Nothing
659 -- > lookupLT 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
660 lookupLT :: Ord k => k -> Map k v -> Maybe (k, v)
661 lookupLT = goNothing
662 where
663 goNothing !_ Tip = Nothing
664 goNothing k (Bin _ kx x l r) | k <= kx = goNothing k l
665 | otherwise = goJust k kx x r
666
667 goJust !_ kx' x' Tip = Just (kx', x')
668 goJust k kx' x' (Bin _ kx x l r) | k <= kx = goJust k kx' x' l
669 | otherwise = goJust k kx x r
670 #if __GLASGOW_HASKELL__
671 {-# INLINABLE lookupLT #-}
672 #else
673 {-# INLINE lookupLT #-}
674 #endif
675
676 -- | /O(log n)/. Find smallest key greater than the given one and return the
677 -- corresponding (key, value) pair.
678 --
679 -- > lookupGT 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
680 -- > lookupGT 5 (fromList [(3,'a'), (5,'b')]) == Nothing
681 lookupGT :: Ord k => k -> Map k v -> Maybe (k, v)
682 lookupGT = goNothing
683 where
684 goNothing !_ Tip = Nothing
685 goNothing k (Bin _ kx x l r) | k < kx = goJust k kx x l
686 | otherwise = goNothing k r
687
688 goJust !_ kx' x' Tip = Just (kx', x')
689 goJust k kx' x' (Bin _ kx x l r) | k < kx = goJust k kx x l
690 | otherwise = goJust k kx' x' r
691 #if __GLASGOW_HASKELL__
692 {-# INLINABLE lookupGT #-}
693 #else
694 {-# INLINE lookupGT #-}
695 #endif
696
697 -- | /O(log n)/. Find largest key smaller or equal to the given one and return
698 -- the corresponding (key, value) pair.
699 --
700 -- > lookupLE 2 (fromList [(3,'a'), (5,'b')]) == Nothing
701 -- > lookupLE 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
702 -- > lookupLE 5 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
703 lookupLE :: Ord k => k -> Map k v -> Maybe (k, v)
704 lookupLE = goNothing
705 where
706 goNothing !_ Tip = Nothing
707 goNothing k (Bin _ kx x l r) = case compare k kx of LT -> goNothing k l
708 EQ -> Just (kx, x)
709 GT -> goJust k kx x r
710
711 goJust !_ kx' x' Tip = Just (kx', x')
712 goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx' x' l
713 EQ -> Just (kx, x)
714 GT -> goJust k kx x r
715 #if __GLASGOW_HASKELL__
716 {-# INLINABLE lookupLE #-}
717 #else
718 {-# INLINE lookupLE #-}
719 #endif
720
721 -- | /O(log n)/. Find smallest key greater or equal to the given one and return
722 -- the corresponding (key, value) pair.
723 --
724 -- > lookupGE 3 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a')
725 -- > lookupGE 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b')
726 -- > lookupGE 6 (fromList [(3,'a'), (5,'b')]) == Nothing
727 lookupGE :: Ord k => k -> Map k v -> Maybe (k, v)
728 lookupGE = goNothing
729 where
730 goNothing !_ Tip = Nothing
731 goNothing k (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx x l
732 EQ -> Just (kx, x)
733 GT -> goNothing k r
734
735 goJust !_ kx' x' Tip = Just (kx', x')
736 goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx x l
737 EQ -> Just (kx, x)
738 GT -> goJust k kx' x' r
739 #if __GLASGOW_HASKELL__
740 {-# INLINABLE lookupGE #-}
741 #else
742 {-# INLINE lookupGE #-}
743 #endif
744
745 {--------------------------------------------------------------------
746 Construction
747 --------------------------------------------------------------------}
748 -- | /O(1)/. The empty map.
749 --
750 -- > empty == fromList []
751 -- > size empty == 0
752
753 empty :: Map k a
754 empty = Tip
755 {-# INLINE empty #-}
756
757 -- | /O(1)/. A map with a single element.
758 --
759 -- > singleton 1 'a' == fromList [(1, 'a')]
760 -- > size (singleton 1 'a') == 1
761
762 singleton :: k -> a -> Map k a
763 singleton k x = Bin 1 k x Tip Tip
764 {-# INLINE singleton #-}
765
766 {--------------------------------------------------------------------
767 Insertion
768 --------------------------------------------------------------------}
769 -- | /O(log n)/. Insert a new key and value in the map.
770 -- If the key is already present in the map, the associated value is
771 -- replaced with the supplied value. 'insert' is equivalent to
772 -- @'insertWith' 'const'@.
773 --
774 -- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')]
775 -- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')]
776 -- > insert 5 'x' empty == singleton 5 'x'
777
778 -- See Note: Type of local 'go' function
779 -- See Note: Avoiding worker/wrapper
780 insert :: Ord k => k -> a -> Map k a -> Map k a
781 insert kx0 = go kx0 kx0
782 where
783 -- Unlike insertR, we only get sharing here
784 -- when the inserted value is at the same address
785 -- as the present value. We try anyway; this condition
786 -- seems particularly likely to occur in 'union'.
787 go :: Ord k => k -> k -> a -> Map k a -> Map k a
788 go orig !_ x Tip = singleton (lazy orig) x
789 go orig !kx x t@(Bin sz ky y l r) =
790 case compare kx ky of
791 LT | l' `ptrEq` l -> t
792 | otherwise -> balanceL ky y l' r
793 where !l' = go orig kx x l
794 GT | r' `ptrEq` r -> t
795 | otherwise -> balanceR ky y l r'
796 where !r' = go orig kx x r
797 EQ | x `ptrEq` y && (lazy orig `seq` (orig `ptrEq` ky)) -> t
798 | otherwise -> Bin sz (lazy orig) x l r
799 #if __GLASGOW_HASKELL__
800 {-# INLINABLE insert #-}
801 #else
802 {-# INLINE insert #-}
803 #endif
804
805 #ifndef __GLASGOW_HASKELL__
806 lazy :: a -> a
807 lazy a = a
808 #endif
809
810 -- [Note: Avoiding worker/wrapper]
811 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
812 -- 'insert' has to go to great lengths to get pointer equality right and
813 -- to prevent unnecessary allocation. The trouble is that GHC *really* wants
814 -- to unbox the key and throw away the boxed one. This is bad for us, because
815 -- we want to compare the pointer of the box we are given to the one already
816 -- present if they compare EQ. It's also bad for us because it leads to the
817 -- key being *reboxed* if it's actually stored in the map. Ugh! So we pass the
818 -- 'go' function *two copies* of the key we're given. One of them we use for
819 -- comparisons; the other we keep in our pocket. To prevent worker/wrapper from
820 -- messing with the copy in our pocket, we sprinkle about calls to the magical
821 -- function 'lazy'. This is all horrible, but it seems to work okay.
822
823
824 -- Insert a new key and value in the map if it is not already present.
825 -- Used by `union`.
826
827 -- See Note: Type of local 'go' function
828 -- See Note: Avoiding worker/wrapper
829 insertR :: Ord k => k -> a -> Map k a -> Map k a
830 insertR kx0 = go kx0 kx0
831 where
832 go :: Ord k => k -> k -> a -> Map k a -> Map k a
833 go orig !_ x Tip = singleton (lazy orig) x
834 go orig !kx x t@(Bin _ ky y l r) =
835 case compare kx ky of
836 LT | l' `ptrEq` l -> t
837 | otherwise -> balanceL ky y l' r
838 where !l' = go orig kx x l
839 GT | r' `ptrEq` r -> t
840 | otherwise -> balanceR ky y l r'
841 where !r' = go orig kx x r
842 EQ -> t
843 #if __GLASGOW_HASKELL__
844 {-# INLINABLE insertR #-}
845 #else
846 {-# INLINE insertR #-}
847 #endif
848
849 -- | /O(log n)/. Insert with a function, combining new value and old value.
850 -- @'insertWith' f key value mp@
851 -- will insert the pair (key, value) into @mp@ if key does
852 -- not exist in the map. If the key does exist, the function will
853 -- insert the pair @(key, f new_value old_value)@.
854 --
855 -- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")]
856 -- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
857 -- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx"
858
859 insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
860 insertWith = go
861 where
862 -- We have no hope of making pointer equality tricks work
863 -- here, because lazy insertWith *always* changes the tree,
864 -- either adding a new entry or replacing an element with a
865 -- thunk.
866 go :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
867 go _ !kx x Tip = singleton kx x
868 go f !kx x (Bin sy ky y l r) =
869 case compare kx ky of
870 LT -> balanceL ky y (go f kx x l) r
871 GT -> balanceR ky y l (go f kx x r)
872 EQ -> Bin sy kx (f x y) l r
873
874 #if __GLASGOW_HASKELL__
875 {-# INLINABLE insertWith #-}
876 #else
877 {-# INLINE insertWith #-}
878 #endif
879
880 -- | A helper function for 'unionWith'. When the key is already in
881 -- the map, the key is left alone, not replaced. The combining
882 -- function is flipped--it is applied to the old value and then the
883 -- new value.
884
885 insertWithR :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
886 insertWithR = go
887 where
888 go :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
889 go _ !kx x Tip = singleton kx x
890 go f !kx x (Bin sy ky y l r) =
891 case compare kx ky of
892 LT -> balanceL ky y (go f kx x l) r
893 GT -> balanceR ky y l (go f kx x r)
894 EQ -> Bin sy ky (f y x) l r
895 #if __GLASGOW_HASKELL__
896 {-# INLINABLE insertWithR #-}
897 #else
898 {-# INLINE insertWithR #-}
899 #endif
900
901 -- | /O(log n)/. Insert with a function, combining key, new value and old value.
902 -- @'insertWithKey' f key value mp@
903 -- will insert the pair (key, value) into @mp@ if key does
904 -- not exist in the map. If the key does exist, the function will
905 -- insert the pair @(key,f key new_value old_value)@.
906 -- Note that the key passed to f is the same key passed to 'insertWithKey'.
907 --
908 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
909 -- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")]
910 -- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
911 -- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx"
912
913 -- See Note: Type of local 'go' function
914 insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
915 insertWithKey = go
916 where
917 go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
918 go _ !kx x Tip = singleton kx x
919 go f kx x (Bin sy ky y l r) =
920 case compare kx ky of
921 LT -> balanceL ky y (go f kx x l) r
922 GT -> balanceR ky y l (go f kx x r)
923 EQ -> Bin sy kx (f kx x y) l r
924 #if __GLASGOW_HASKELL__
925 {-# INLINABLE insertWithKey #-}
926 #else
927 {-# INLINE insertWithKey #-}
928 #endif
929
930 -- | A helper function for 'unionWithKey'. When the key is already in
931 -- the map, the key is left alone, not replaced. The combining
932 -- function is flipped--it is applied to the old value and then the
933 -- new value.
934 insertWithKeyR :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
935 insertWithKeyR = go
936 where
937 go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
938 go _ !kx x Tip = singleton kx x
939 go f kx x (Bin sy ky y l r) =
940 case compare kx ky of
941 LT -> balanceL ky y (go f kx x l) r
942 GT -> balanceR ky y l (go f kx x r)
943 EQ -> Bin sy ky (f ky y x) l r
944 #if __GLASGOW_HASKELL__
945 {-# INLINABLE insertWithKeyR #-}
946 #else
947 {-# INLINE insertWithKeyR #-}
948 #endif
949
950 -- | /O(log n)/. Combines insert operation with old value retrieval.
951 -- The expression (@'insertLookupWithKey' f k x map@)
952 -- is a pair where the first element is equal to (@'lookup' k map@)
953 -- and the second element equal to (@'insertWithKey' f k x map@).
954 --
955 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
956 -- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")])
957 -- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")])
958 -- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx")
959 --
960 -- This is how to define @insertLookup@ using @insertLookupWithKey@:
961 --
962 -- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t
963 -- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")])
964 -- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")])
965
966 -- See Note: Type of local 'go' function
967 insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a
968 -> (Maybe a, Map k a)
969 insertLookupWithKey f0 k0 x0 = toPair . go f0 k0 x0
970 where
971 go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> StrictPair (Maybe a) (Map k a)
972 go _ !kx x Tip = (Nothing :*: singleton kx x)
973 go f kx x (Bin sy ky y l r) =
974 case compare kx ky of
975 LT -> let !(found :*: l') = go f kx x l
976 !t' = balanceL ky y l' r
977 in (found :*: t')
978 GT -> let !(found :*: r') = go f kx x r
979 !t' = balanceR ky y l r'
980 in (found :*: t')
981 EQ -> (Just y :*: Bin sy kx (f kx x y) l r)
982 #if __GLASGOW_HASKELL__
983 {-# INLINABLE insertLookupWithKey #-}
984 #else
985 {-# INLINE insertLookupWithKey #-}
986 #endif
987
988 {--------------------------------------------------------------------
989 Deletion
990 --------------------------------------------------------------------}
991 -- | /O(log n)/. Delete a key and its value from the map. When the key is not
992 -- a member of the map, the original map is returned.
993 --
994 -- > delete 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
995 -- > delete 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
996 -- > delete 5 empty == empty
997
998 -- See Note: Type of local 'go' function
999 delete :: Ord k => k -> Map k a -> Map k a
1000 delete = go
1001 where
1002 go :: Ord k => k -> Map k a -> Map k a
1003 go !_ Tip = Tip
1004 go k t@(Bin _ kx x l r) =
1005 case compare k kx of
1006 LT | l' `ptrEq` l -> t
1007 | otherwise -> balanceR kx x l' r
1008 where !l' = go k l
1009 GT | r' `ptrEq` r -> t
1010 | otherwise -> balanceL kx x l r'
1011 where !r' = go k r
1012 EQ -> glue l r
1013 #if __GLASGOW_HASKELL__
1014 {-# INLINABLE delete #-}
1015 #else
1016 {-# INLINE delete #-}
1017 #endif
1018
1019 -- | /O(log n)/. Update a value at a specific key with the result of the provided function.
1020 -- When the key is not
1021 -- a member of the map, the original map is returned.
1022 --
1023 -- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
1024 -- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
1025 -- > adjust ("new " ++) 7 empty == empty
1026
1027 adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a
1028 adjust f = adjustWithKey (\_ x -> f x)
1029 #if __GLASGOW_HASKELL__
1030 {-# INLINABLE adjust #-}
1031 #else
1032 {-# INLINE adjust #-}
1033 #endif
1034
1035 -- | /O(log n)/. Adjust a value at a specific key. When the key is not
1036 -- a member of the map, the original map is returned.
1037 --
1038 -- > let f key x = (show key) ++ ":new " ++ x
1039 -- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
1040 -- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
1041 -- > adjustWithKey f 7 empty == empty
1042
1043 adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a
1044 adjustWithKey = go
1045 where
1046 go :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a
1047 go _ !_ Tip = Tip
1048 go f k (Bin sx kx x l r) =
1049 case compare k kx of
1050 LT -> Bin sx kx x (go f k l) r
1051 GT -> Bin sx kx x l (go f k r)
1052 EQ -> Bin sx kx (f kx x) l r
1053 #if __GLASGOW_HASKELL__
1054 {-# INLINABLE adjustWithKey #-}
1055 #else
1056 {-# INLINE adjustWithKey #-}
1057 #endif
1058
1059 -- | /O(log n)/. The expression (@'update' f k map@) updates the value @x@
1060 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
1061 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
1062 --
1063 -- > let f x = if x == "a" then Just "new a" else Nothing
1064 -- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
1065 -- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
1066 -- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1067
1068 update :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
1069 update f = updateWithKey (\_ x -> f x)
1070 #if __GLASGOW_HASKELL__
1071 {-# INLINABLE update #-}
1072 #else
1073 {-# INLINE update #-}
1074 #endif
1075
1076 -- | /O(log n)/. The expression (@'updateWithKey' f k map@) updates the
1077 -- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing',
1078 -- the element is deleted. If it is (@'Just' y@), the key @k@ is bound
1079 -- to the new value @y@.
1080 --
1081 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
1082 -- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
1083 -- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
1084 -- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1085
1086 -- See Note: Type of local 'go' function
1087 updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a
1088 updateWithKey = go
1089 where
1090 go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a
1091 go _ !_ Tip = Tip
1092 go f k(Bin sx kx x l r) =
1093 case compare k kx of
1094 LT -> balanceR kx x (go f k l) r
1095 GT -> balanceL kx x l (go f k r)
1096 EQ -> case f kx x of
1097 Just x' -> Bin sx kx x' l r
1098 Nothing -> glue l r
1099 #if __GLASGOW_HASKELL__
1100 {-# INLINABLE updateWithKey #-}
1101 #else
1102 {-# INLINE updateWithKey #-}
1103 #endif
1104
1105 -- | /O(log n)/. Lookup and update. See also 'updateWithKey'.
1106 -- The function returns changed value, if it is updated.
1107 -- Returns the original key value if the map entry is deleted.
1108 --
1109 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
1110 -- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "5:new a", fromList [(3, "b"), (5, "5:new a")])
1111 -- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")])
1112 -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a")
1113
1114 -- See Note: Type of local 'go' function
1115 updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a)
1116 updateLookupWithKey f0 k0 = toPair . go f0 k0
1117 where
1118 go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> StrictPair (Maybe a) (Map k a)
1119 go _ !_ Tip = (Nothing :*: Tip)
1120 go f k (Bin sx kx x l r) =
1121 case compare k kx of
1122 LT -> let !(found :*: l') = go f k l
1123 !t' = balanceR kx x l' r
1124 in (found :*: t')
1125 GT -> let !(found :*: r') = go f k r
1126 !t' = balanceL kx x l r'
1127 in (found :*: t')
1128 EQ -> case f kx x of
1129 Just x' -> (Just x' :*: Bin sx kx x' l r)
1130 Nothing -> let !glued = glue l r
1131 in (Just x :*: glued)
1132 #if __GLASGOW_HASKELL__
1133 {-# INLINABLE updateLookupWithKey #-}
1134 #else
1135 {-# INLINE updateLookupWithKey #-}
1136 #endif
1137
1138 -- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
1139 -- 'alter' can be used to insert, delete, or update a value in a 'Map'.
1140 -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
1141 --
1142 -- > let f _ = Nothing
1143 -- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
1144 -- > alter f 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1145 -- >
1146 -- > let f _ = Just "c"
1147 -- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "c")]
1148 -- > alter f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "c")]
1149
1150 -- See Note: Type of local 'go' function
1151 alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
1152 alter = go
1153 where
1154 go :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
1155 go f !k Tip = case f Nothing of
1156 Nothing -> Tip
1157 Just x -> singleton k x
1158
1159 go f k (Bin sx kx x l r) = case compare k kx of
1160 LT -> balance kx x (go f k l) r
1161 GT -> balance kx x l (go f k r)
1162 EQ -> case f (Just x) of
1163 Just x' -> Bin sx kx x' l r
1164 Nothing -> glue l r
1165 #if __GLASGOW_HASKELL__
1166 {-# INLINABLE alter #-}
1167 #else
1168 {-# INLINE alter #-}
1169 #endif
1170
1171 -- Used to choose the appropriate alterF implementation.
1172 data AreWeStrict = Strict | Lazy
1173
1174 -- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at
1175 -- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete,
1176 -- or update a value in a 'Map'. In short: @'lookup' k \<$\> 'alterF' f k m = f
1177 -- ('lookup' k m)@.
1178 --
1179 -- Example:
1180 --
1181 -- @
1182 -- interactiveAlter :: Int -> Map Int String -> IO (Map Int String)
1183 -- interactiveAlter k m = alterF f k m where
1184 -- f Nothing = do
1185 -- putStrLn $ show k ++
1186 -- " was not found in the map. Would you like to add it?"
1187 -- getUserResponse1 :: IO (Maybe String)
1188 -- f (Just old) = do
1189 -- putStrLn $ "The key is currently bound to " ++ show old ++
1190 -- ". Would you like to change or delete it?"
1191 -- getUserResponse2 :: IO (Maybe String)
1192 -- @
1193 --
1194 -- 'alterF' is the most general operation for working with an individual
1195 -- key that may or may not be in a given map. When used with trivial
1196 -- functors like 'Identity' and 'Const', it is often slightly slower than
1197 -- more specialized combinators like 'lookup' and 'insert'. However, when
1198 -- the functor is non-trivial and key comparison is not particularly cheap,
1199 -- it is the fastest way.
1200 --
1201 -- Note on rewrite rules:
1202 --
1203 -- This module includes GHC rewrite rules to optimize 'alterF' for
1204 -- the 'Const' and 'Identity' functors. In general, these rules
1205 -- improve performance. The sole exception is that when using
1206 -- 'Identity', deleting a key that is already absent takes longer
1207 -- than it would without the rules. If you expect this to occur
1208 -- a very large fraction of the time, you might consider using a
1209 -- private copy of the 'Identity' type.
1210 --
1211 -- Note: 'alterF' is a flipped version of the @at@ combinator from
1212 -- @Control.Lens.At@.
1213 --
1214 -- @since 0.5.8
1215 alterF :: (Functor f, Ord k)
1216 => (Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
1217 alterF f k m = atKeyImpl Lazy k f m
1218
1219 #ifndef __GLASGOW_HASKELL__
1220 {-# INLINE alterF #-}
1221 #else
1222 {-# INLINABLE [2] alterF #-}
1223
1224 -- We can save a little time by recognizing the special case of
1225 -- `Control.Applicative.Const` and just doing a lookup.
1226 {-# RULES
1227 "alterF/Const" forall k (f :: Maybe a -> Const b (Maybe a)) . alterF f k = \m -> Const . getConst . f $ lookup k m
1228 #-}
1229
1230 #if MIN_VERSION_base(4,8,0)
1231 -- base 4.8 and above include Data.Functor.Identity, so we can
1232 -- save a pretty decent amount of time by handling it specially.
1233 {-# RULES
1234 "alterF/Identity" forall k f . alterF f k = atKeyIdentity k f
1235 #-}
1236 #endif
1237 #endif
1238
1239 atKeyImpl :: (Functor f, Ord k) =>
1240 AreWeStrict -> k -> (Maybe a -> f (Maybe a)) -> Map k a -> f (Map k a)
1241 #ifdef DEFINE_ALTERF_FALLBACK
1242 atKeyImpl strict !k f m
1243 -- It doesn't seem sensible to worry about overflowing the queue
1244 -- if the word size is 61 or more. If I calculate it correctly,
1245 -- that would take a map with nearly a quadrillion entries.
1246 | wordSize < 61 && size m >= alterFCutoff = alterFFallback strict k f m
1247 #endif
1248 atKeyImpl strict !k f m = case lookupTrace k m of
1249 TraceResult mv q -> (<$> f mv) $ \ fres ->
1250 case fres of
1251 Nothing -> case mv of
1252 Nothing -> m
1253 Just old -> deleteAlong old q m
1254 Just new -> case strict of
1255 Strict -> new `seq` case mv of
1256 Nothing -> insertAlong q k new m
1257 Just _ -> replaceAlong q new m
1258 Lazy -> case mv of
1259 Nothing -> insertAlong q k new m
1260 Just _ -> replaceAlong q new m
1261
1262 {-# INLINE atKeyImpl #-}
1263
1264 #ifdef DEFINE_ALTERF_FALLBACK
1265 alterFCutoff :: Int
1266 #if WORD_SIZE_IN_BITS == 32
1267 alterFCutoff = 55744454
1268 #else
1269 alterFCutoff = case wordSize of
1270 30 -> 17637893
1271 31 -> 31356255
1272 32 -> 55744454
1273 x -> (4^(x*2-2)) `quot` (3^(x*2-2)) -- Unlikely
1274 #endif
1275 #endif
1276
1277 data TraceResult a = TraceResult (Maybe a) {-# UNPACK #-} !BitQueue
1278
1279 -- Look up a key and return a result indicating whether it was found
1280 -- and what path was taken.
1281 lookupTrace :: Ord k => k -> Map k a -> TraceResult a
1282 lookupTrace = go emptyQB
1283 where
1284 go :: Ord k => BitQueueB -> k -> Map k a -> TraceResult a
1285 go !q !_ Tip = TraceResult Nothing (buildQ q)
1286 go q k (Bin _ kx x l r) = case compare k kx of
1287 LT -> (go $! q `snocQB` False) k l
1288 GT -> (go $! q `snocQB` True) k r
1289 EQ -> TraceResult (Just x) (buildQ q)
1290
1291 -- GHC 7.8 doesn't manage to unbox the queue properly
1292 -- unless we explicitly inline this function. This stuff
1293 -- is a bit touchy, unfortunately.
1294 #if __GLASGOW_HASKELL__ >= 710
1295 {-# INLINABLE lookupTrace #-}
1296 #else
1297 {-# INLINE lookupTrace #-}
1298 #endif
1299
1300 -- Insert at a location (which will always be a leaf)
1301 -- described by the path passed in.
1302 insertAlong :: BitQueue -> k -> a -> Map k a -> Map k a
1303 insertAlong !_ kx x Tip = singleton kx x
1304 insertAlong q kx x (Bin sz ky y l r) =
1305 case unconsQ q of
1306 Just (False, tl) -> balanceL ky y (insertAlong tl kx x l) r
1307 Just (True,tl) -> balanceR ky y l (insertAlong tl kx x r)
1308 Nothing -> Bin sz kx x l r -- Shouldn't happen
1309
1310 -- Delete from a location (which will always be a node)
1311 -- described by the path passed in.
1312 --
1313 -- This is fairly horrifying! We don't actually have any
1314 -- use for the old value we're deleting. But if GHC sees
1315 -- that, then it will allocate a thunk representing the
1316 -- Map with the key deleted before we have any reason to
1317 -- believe we'll actually want that. This transformation
1318 -- enhances sharing, but we don't care enough about that.
1319 -- So deleteAlong needs to take the old value, and we need
1320 -- to convince GHC somehow that it actually uses it. We
1321 -- can't NOINLINE deleteAlong, because that would prevent
1322 -- the BitQueue from being unboxed. So instead we pass the
1323 -- old value to a NOINLINE constant function and then
1324 -- convince GHC that we use the result throughout the
1325 -- computation. Doing the obvious thing and just passing
1326 -- the value itself through the recursion costs 3-4% time,
1327 -- so instead we convert the value to a magical zero-width
1328 -- proxy that's ultimately erased.
1329 deleteAlong :: any -> BitQueue -> Map k a -> Map k a
1330 deleteAlong old !q0 !m = go (bogus old) q0 m where
1331 #ifdef USE_MAGIC_PROXY
1332 go :: Proxy# () -> BitQueue -> Map k a -> Map k a
1333 #else
1334 go :: any -> BitQueue -> Map k a -> Map k a
1335 #endif
1336 go !_ !_ Tip = Tip
1337 go foom q (Bin _ ky y l r) =
1338 case unconsQ q of
1339 Just (False, tl) -> balanceR ky y (go foom tl l) r
1340 Just (True, tl) -> balanceL ky y l (go foom tl r)
1341 Nothing -> glue l r
1342
1343 #ifdef USE_MAGIC_PROXY
1344 {-# NOINLINE bogus #-}
1345 bogus :: a -> Proxy# ()
1346 bogus _ = proxy#
1347 #else
1348 -- No point hiding in this case.
1349 {-# INLINE bogus #-}
1350 bogus :: a -> a
1351 bogus a = a
1352 #endif
1353
1354 -- Replace the value found in the node described
1355 -- by the given path with a new one.
1356 replaceAlong :: BitQueue -> a -> Map k a -> Map k a
1357 replaceAlong !_ _ Tip = Tip -- Should not happen
1358 replaceAlong q x (Bin sz ky y l r) =
1359 case unconsQ q of
1360 Just (False, tl) -> Bin sz ky y (replaceAlong tl x l) r
1361 Just (True,tl) -> Bin sz ky y l (replaceAlong tl x r)
1362 Nothing -> Bin sz ky x l r
1363
1364 #if __GLASGOW_HASKELL__ && MIN_VERSION_base(4,8,0)
1365 atKeyIdentity :: Ord k => k -> (Maybe a -> Identity (Maybe a)) -> Map k a -> Identity (Map k a)
1366 atKeyIdentity k f t = Identity $ atKeyPlain Lazy k (coerce f) t
1367 {-# INLINABLE atKeyIdentity #-}
1368
1369 atKeyPlain :: Ord k => AreWeStrict -> k -> (Maybe a -> Maybe a) -> Map k a -> Map k a
1370 atKeyPlain strict k0 f0 t = case go k0 f0 t of
1371 AltSmaller t' -> t'
1372 AltBigger t' -> t'
1373 AltAdj t' -> t'
1374 AltSame -> t
1375 where
1376 go :: Ord k => k -> (Maybe a -> Maybe a) -> Map k a -> Altered k a
1377 go !k f Tip = case f Nothing of
1378 Nothing -> AltSame
1379 Just x -> case strict of
1380 Lazy -> AltBigger $ singleton k x
1381 Strict -> x `seq` (AltBigger $ singleton k x)
1382
1383 go k f (Bin sx kx x l r) = case compare k kx of
1384 LT -> case go k f l of
1385 AltSmaller l' -> AltSmaller $ balanceR kx x l' r
1386 AltBigger l' -> AltBigger $ balanceL kx x l' r
1387 AltAdj l' -> AltAdj $ Bin sx kx x l' r
1388 AltSame -> AltSame
1389 GT -> case go k f r of
1390 AltSmaller r' -> AltSmaller $ balanceL kx x l r'
1391 AltBigger r' -> AltBigger $ balanceR kx x l r'
1392 AltAdj r' -> AltAdj $ Bin sx kx x l r'
1393 AltSame -> AltSame
1394 EQ -> case f (Just x) of
1395 Just x' -> case strict of
1396 Lazy -> AltAdj $ Bin sx kx x' l r
1397 Strict -> x' `seq` (AltAdj $ Bin sx kx x' l r)
1398 Nothing -> AltSmaller $ glue l r
1399 {-# INLINE atKeyPlain #-}
1400
1401 data Altered k a = AltSmaller !(Map k a) | AltBigger !(Map k a) | AltAdj !(Map k a) | AltSame
1402 #endif
1403
1404 #ifdef DEFINE_ALTERF_FALLBACK
1405 -- When the map is too large to use a bit queue, we fall back to
1406 -- this much slower version which uses a more "natural" implementation
1407 -- improved with Yoneda to avoid repeated fmaps. This works okayish for
1408 -- some operations, but it's pretty lousy for lookups.
1409 alterFFallback :: (Functor f, Ord k)
1410 => AreWeStrict -> k -> (Maybe a -> f (Maybe a)) -> Map k a -> f (Map k a)
1411 alterFFallback Lazy k f t = alterFYoneda k (\m q -> q <$> f m) t id
1412 alterFFallback Strict k f t = alterFYoneda k (\m q -> q . forceMaybe <$> f m) t id
1413 where
1414 forceMaybe Nothing = Nothing
1415 forceMaybe may@(Just !_) = may
1416 {-# NOINLINE alterFFallback #-}
1417
1418 alterFYoneda :: Ord k =>
1419 k -> (Maybe a -> (Maybe a -> b) -> f b) -> Map k a -> (Map k a -> b) -> f b
1420 alterFYoneda = go
1421 where
1422 go :: Ord k =>
1423 k -> (Maybe a -> (Maybe a -> b) -> f b) -> Map k a -> (Map k a -> b) -> f b
1424 go !k f Tip g = f Nothing $ \ mx -> case mx of
1425 Nothing -> g Tip
1426 Just x -> g (singleton k x)
1427 go k f (Bin sx kx x l r) g = case compare k kx of
1428 LT -> go k f l (\m -> g (balance kx x m r))
1429 GT -> go k f r (\m -> g (balance kx x l m))
1430 EQ -> f (Just x) $ \ mx' -> case mx' of
1431 Just x' -> g (Bin sx kx x' l r)
1432 Nothing -> g (glue l r)
1433 {-# INLINE alterFYoneda #-}
1434 #endif
1435
1436 {--------------------------------------------------------------------
1437 Indexing
1438 --------------------------------------------------------------------}
1439 -- | /O(log n)/. Return the /index/ of a key, which is its zero-based index in
1440 -- the sequence sorted by keys. The index is a number from /0/ up to, but not
1441 -- including, the 'size' of the map. Calls 'error' when the key is not
1442 -- a 'member' of the map.
1443 --
1444 -- > findIndex 2 (fromList [(5,"a"), (3,"b")]) Error: element is not in the map
1445 -- > findIndex 3 (fromList [(5,"a"), (3,"b")]) == 0
1446 -- > findIndex 5 (fromList [(5,"a"), (3,"b")]) == 1
1447 -- > findIndex 6 (fromList [(5,"a"), (3,"b")]) Error: element is not in the map
1448
1449 -- See Note: Type of local 'go' function
1450 findIndex :: Ord k => k -> Map k a -> Int
1451 findIndex = go 0
1452 where
1453 go :: Ord k => Int -> k -> Map k a -> Int
1454 go !_ !_ Tip = error "Map.findIndex: element is not in the map"
1455 go idx k (Bin _ kx _ l r) = case compare k kx of
1456 LT -> go idx k l
1457 GT -> go (idx + size l + 1) k r
1458 EQ -> idx + size l
1459 #if __GLASGOW_HASKELL__
1460 {-# INLINABLE findIndex #-}
1461 #endif
1462
1463 -- | /O(log n)/. Lookup the /index/ of a key, which is its zero-based index in
1464 -- the sequence sorted by keys. The index is a number from /0/ up to, but not
1465 -- including, the 'size' of the map.
1466 --
1467 -- > isJust (lookupIndex 2 (fromList [(5,"a"), (3,"b")])) == False
1468 -- > fromJust (lookupIndex 3 (fromList [(5,"a"), (3,"b")])) == 0
1469 -- > fromJust (lookupIndex 5 (fromList [(5,"a"), (3,"b")])) == 1
1470 -- > isJust (lookupIndex 6 (fromList [(5,"a"), (3,"b")])) == False
1471
1472 -- See Note: Type of local 'go' function
1473 lookupIndex :: Ord k => k -> Map k a -> Maybe Int
1474 lookupIndex = go 0
1475 where
1476 go :: Ord k => Int -> k -> Map k a -> Maybe Int
1477 go !_ !_ Tip = Nothing
1478 go idx k (Bin _ kx _ l r) = case compare k kx of
1479 LT -> go idx k l
1480 GT -> go (idx + size l + 1) k r
1481 EQ -> Just $! idx + size l
1482 #if __GLASGOW_HASKELL__
1483 {-# INLINABLE lookupIndex #-}
1484 #endif
1485
1486 -- | /O(log n)/. Retrieve an element by its /index/, i.e. by its zero-based
1487 -- index in the sequence sorted by keys. If the /index/ is out of range (less
1488 -- than zero, greater or equal to 'size' of the map), 'error' is called.
1489 --
1490 -- > elemAt 0 (fromList [(5,"a"), (3,"b")]) == (3,"b")
1491 -- > elemAt 1 (fromList [(5,"a"), (3,"b")]) == (5, "a")
1492 -- > elemAt 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range
1493
1494 elemAt :: Int -> Map k a -> (k,a)
1495 elemAt !_ Tip = error "Map.elemAt: index out of range"
1496 elemAt i (Bin _ kx x l r)
1497 = case compare i sizeL of
1498 LT -> elemAt i l
1499 GT -> elemAt (i-sizeL-1) r
1500 EQ -> (kx,x)
1501 where
1502 sizeL = size l
1503
1504 -- | Take a given number of entries in key order, beginning
1505 -- with the smallest keys.
1506 --
1507 -- @
1508 -- take n = 'fromDistinctAscList' . 'Prelude.take' n . 'toAscList'
1509 -- @
1510 --
1511 -- @since 0.5.8
1512
1513 take :: Int -> Map k a -> Map k a
1514 take i m | i >= size m = m
1515 take i0 m0 = go i0 m0
1516 where
1517 go i !_ | i <= 0 = Tip
1518 go !_ Tip = Tip
1519 go i (Bin _ kx x l r) =
1520 case compare i sizeL of
1521 LT -> go i l
1522 GT -> link kx x l (go (i - sizeL - 1) r)
1523 EQ -> l
1524 where sizeL = size l
1525
1526 -- | Drop a given number of entries in key order, beginning
1527 -- with the smallest keys.
1528 --
1529 -- @
1530 -- drop n = 'fromDistinctAscList' . 'Prelude.drop' n . 'toAscList'
1531 -- @
1532 --
1533 -- @since 0.5.8
1534 drop :: Int -> Map k a -> Map k a
1535 drop i m | i >= size m = Tip
1536 drop i0 m0 = go i0 m0
1537 where
1538 go i m | i <= 0 = m
1539 go !_ Tip = Tip
1540 go i (Bin _ kx x l r) =
1541 case compare i sizeL of
1542 LT -> link kx x (go i l) r
1543 GT -> go (i - sizeL - 1) r
1544 EQ -> insertMin kx x r
1545 where sizeL = size l
1546
1547 -- | /O(log n)/. Split a map at a particular index.
1548 --
1549 -- @
1550 -- splitAt !n !xs = ('take' n xs, 'drop' n xs)
1551 -- @
1552 --
1553 -- @since 0.5.8
1554 splitAt :: Int -> Map k a -> (Map k a, Map k a)
1555 splitAt i0 m0
1556 | i0 >= size m0 = (m0, Tip)
1557 | otherwise = toPair $ go i0 m0
1558 where
1559 go i m | i <= 0 = Tip :*: m
1560 go !_ Tip = Tip :*: Tip
1561 go i (Bin _ kx x l r)
1562 = case compare i sizeL of
1563 LT -> case go i l of
1564 ll :*: lr -> ll :*: link kx x lr r
1565 GT -> case go (i - sizeL - 1) r of
1566 rl :*: rr -> link kx x l rl :*: rr
1567 EQ -> l :*: insertMin kx x r
1568 where sizeL = size l
1569
1570 -- | /O(log n)/. Update the element at /index/, i.e. by its zero-based index in
1571 -- the sequence sorted by keys. If the /index/ is out of range (less than zero,
1572 -- greater or equal to 'size' of the map), 'error' is called.
1573 --
1574 -- > updateAt (\ _ _ -> Just "x") 0 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "x"), (5, "a")]
1575 -- > updateAt (\ _ _ -> Just "x") 1 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "x")]
1576 -- > updateAt (\ _ _ -> Just "x") 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range
1577 -- > updateAt (\ _ _ -> Just "x") (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range
1578 -- > updateAt (\_ _ -> Nothing) 0 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1579 -- > updateAt (\_ _ -> Nothing) 1 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1580 -- > updateAt (\_ _ -> Nothing) 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range
1581 -- > updateAt (\_ _ -> Nothing) (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range
1582
1583 updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
1584 updateAt f !i t =
1585 case t of
1586 Tip -> error "Map.updateAt: index out of range"
1587 Bin sx kx x l r -> case compare i sizeL of
1588 LT -> balanceR kx x (updateAt f i l) r
1589 GT -> balanceL kx x l (updateAt f (i-sizeL-1) r)
1590 EQ -> case f kx x of
1591 Just x' -> Bin sx kx x' l r
1592 Nothing -> glue l r
1593 where
1594 sizeL = size l
1595
1596 -- | /O(log n)/. Delete the element at /index/, i.e. by its zero-based index in
1597 -- the sequence sorted by keys. If the /index/ is out of range (less than zero,
1598 -- greater or equal to 'size' of the map), 'error' is called.
1599 --
1600 -- > deleteAt 0 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1601 -- > deleteAt 1 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1602 -- > deleteAt 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range
1603 -- > deleteAt (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range
1604
1605 deleteAt :: Int -> Map k a -> Map k a
1606 deleteAt !i t =
1607 case t of
1608 Tip -> error "Map.deleteAt: index out of range"
1609 Bin _ kx x l r -> case compare i sizeL of
1610 LT -> balanceR kx x (deleteAt i l) r
1611 GT -> balanceL kx x l (deleteAt (i-sizeL-1) r)
1612 EQ -> glue l r
1613 where
1614 sizeL = size l
1615
1616
1617 {--------------------------------------------------------------------
1618 Minimal, Maximal
1619 --------------------------------------------------------------------}
1620
1621 lookupMinSure :: k -> a -> Map k a -> (k, a)
1622 lookupMinSure k a Tip = (k, a)
1623 lookupMinSure _ _ (Bin _ k a l _) = lookupMinSure k a l
1624
1625 -- | /O(log n)/. The minimal key of the map. Returns 'Nothing' if the map is empty.
1626 --
1627 -- > lookupMin (fromList [(5,"a"), (3,"b")]) == Just (3,"b")
1628 -- > findMin empty = Nothing
1629 --
1630 -- @since 0.5.9
1631
1632 lookupMin :: Map k a -> Maybe (k,a)
1633 lookupMin Tip = Nothing
1634 lookupMin (Bin _ k x l _) = Just $! lookupMinSure k x l
1635
1636 -- | /O(log n)/. The minimal key of the map. Calls 'error' if the map is empty.
1637 --
1638 -- > findMin (fromList [(5,"a"), (3,"b")]) == (3,"b")
1639 -- > findMin empty Error: empty map has no minimal element
1640
1641 findMin :: Map k a -> (k,a)
1642 findMin t
1643 | Just r <- lookupMin t = r
1644 | otherwise = error "Map.findMin: empty map has no minimal element"
1645
1646 -- | /O(log n)/. The maximal key of the map. Calls 'error' if the map is empty.
1647 --
1648 -- > findMax (fromList [(5,"a"), (3,"b")]) == (5,"a")
1649 -- > findMax empty Error: empty map has no maximal element
1650
1651 lookupMaxSure :: k -> a -> Map k a -> (k, a)
1652 lookupMaxSure k a Tip = (k, a)
1653 lookupMaxSure _ _ (Bin _ k a _ r) = lookupMaxSure k a r
1654
1655 -- | /O(log n)/. The maximal key of the map. Returns 'Nothing' if the map is empty.
1656 --
1657 -- > lookupMax (fromList [(5,"a"), (3,"b")]) == Just (5,"a")
1658 -- > lookupMax empty = Nothing
1659 --
1660 -- @since 0.5.9
1661
1662 lookupMax :: Map k a -> Maybe (k, a)
1663 lookupMax Tip = Nothing
1664 lookupMax (Bin _ k x _ r) = Just $! lookupMaxSure k x r
1665
1666 findMax :: Map k a -> (k,a)
1667 findMax t
1668 | Just r <- lookupMax t = r
1669 | otherwise = error "Map.findMax: empty map has no maximal element"
1670
1671 -- | /O(log n)/. Delete the minimal key. Returns an empty map if the map is empty.
1672 --
1673 -- > deleteMin (fromList [(5,"a"), (3,"b"), (7,"c")]) == fromList [(5,"a"), (7,"c")]
1674 -- > deleteMin empty == empty
1675
1676 deleteMin :: Map k a -> Map k a
1677 deleteMin (Bin _ _ _ Tip r) = r
1678 deleteMin (Bin _ kx x l r) = balanceR kx x (deleteMin l) r
1679 deleteMin Tip = Tip
1680
1681 -- | /O(log n)/. Delete the maximal key. Returns an empty map if the map is empty.
1682 --
1683 -- > deleteMax (fromList [(5,"a"), (3,"b"), (7,"c")]) == fromList [(3,"b"), (5,"a")]
1684 -- > deleteMax empty == empty
1685
1686 deleteMax :: Map k a -> Map k a
1687 deleteMax (Bin _ _ _ l Tip) = l
1688 deleteMax (Bin _ kx x l r) = balanceL kx x l (deleteMax r)
1689 deleteMax Tip = Tip
1690
1691 -- | /O(log n)/. Update the value at the minimal key.
1692 --
1693 -- > updateMin (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "Xb"), (5, "a")]
1694 -- > updateMin (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1695
1696 updateMin :: (a -> Maybe a) -> Map k a -> Map k a
1697 updateMin f m
1698 = updateMinWithKey (\_ x -> f x) m
1699
1700 -- | /O(log n)/. Update the value at the maximal key.
1701 --
1702 -- > updateMax (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "Xa")]
1703 -- > updateMax (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1704
1705 updateMax :: (a -> Maybe a) -> Map k a -> Map k a
1706 updateMax f m
1707 = updateMaxWithKey (\_ x -> f x) m
1708
1709
1710 -- | /O(log n)/. Update the value at the minimal key.
1711 --
1712 -- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"3:b"), (5,"a")]
1713 -- > updateMinWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1714
1715 updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a
1716 updateMinWithKey _ Tip = Tip
1717 updateMinWithKey f (Bin sx kx x Tip r) = case f kx x of
1718 Nothing -> r
1719 Just x' -> Bin sx kx x' Tip r
1720 updateMinWithKey f (Bin _ kx x l r) = balanceR kx x (updateMinWithKey f l) r
1721
1722 -- | /O(log n)/. Update the value at the maximal key.
1723 --
1724 -- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")]
1725 -- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1726
1727 updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a
1728 updateMaxWithKey _ Tip = Tip
1729 updateMaxWithKey f (Bin sx kx x l Tip) = case f kx x of
1730 Nothing -> l
1731 Just x' -> Bin sx kx x' l Tip
1732 updateMaxWithKey f (Bin _ kx x l r) = balanceL kx x l (updateMaxWithKey f r)
1733
1734 -- | /O(log n)/. Retrieves the minimal (key,value) pair of the map, and
1735 -- the map stripped of that element, or 'Nothing' if passed an empty map.
1736 --
1737 -- > minViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((3,"b"), singleton 5 "a")
1738 -- > minViewWithKey empty == Nothing
1739
1740 minViewWithKey :: Map k a -> Maybe ((k,a), Map k a)
1741 minViewWithKey Tip = Nothing
1742 minViewWithKey (Bin _ k x l r) = Just $
1743 case minViewSure k x l r of
1744 MinView km xm t -> ((km, xm), t)
1745 -- We inline this to give GHC the best possible chance of getting
1746 -- rid of the Maybe and pair constructors, as well as the thunk under
1747 -- the Just.
1748 {-# INLINE minViewWithKey #-}
1749
1750 -- | /O(log n)/. Retrieves the maximal (key,value) pair of the map, and
1751 -- the map stripped of that element, or 'Nothing' if passed an empty map.
1752 --
1753 -- > maxViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((5,"a"), singleton 3 "b")
1754 -- > maxViewWithKey empty == Nothing
1755
1756 maxViewWithKey :: Map k a -> Maybe ((k,a), Map k a)
1757 maxViewWithKey Tip = Nothing
1758 maxViewWithKey (Bin _ k x l r) = Just $
1759 case maxViewSure k x l r of
1760 MaxView km xm t -> ((km, xm), t)
1761 -- See note on inlining at minViewWithKey
1762 {-# INLINE maxViewWithKey #-}
1763
1764 -- | /O(log n)/. Retrieves the value associated with minimal key of the
1765 -- map, and the map stripped of that element, or 'Nothing' if passed an
1766 -- empty map.
1767 --
1768 -- > minView (fromList [(5,"a"), (3,"b")]) == Just ("b", singleton 5 "a")
1769 -- > minView empty == Nothing
1770
1771 minView :: Map k a -> Maybe (a, Map k a)
1772 minView t = case minViewWithKey t of
1773 Nothing -> Nothing
1774 Just ~((_, x), t') -> Just (x, t')
1775
1776 -- | /O(log n)/. Retrieves the value associated with maximal key of the
1777 -- map, and the map stripped of that element, or 'Nothing' if passed an
1778 -- empty map.
1779 --
1780 -- > maxView (fromList [(5,"a"), (3,"b")]) == Just ("a", singleton 3 "b")
1781 -- > maxView empty == Nothing
1782
1783 maxView :: Map k a -> Maybe (a, Map k a)
1784 maxView t = case maxViewWithKey t of
1785 Nothing -> Nothing
1786 Just ~((_, x), t') -> Just (x, t')
1787
1788 {--------------------------------------------------------------------
1789 Union.
1790 --------------------------------------------------------------------}
1791 -- | The union of a list of maps:
1792 -- (@'unions' == 'Prelude.foldl' 'union' 'empty'@).
1793 --
1794 -- > unions [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
1795 -- > == fromList [(3, "b"), (5, "a"), (7, "C")]
1796 -- > unions [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])]
1797 -- > == fromList [(3, "B3"), (5, "A3"), (7, "C")]
1798
1799 unions :: (Foldable f, Ord k) => f (Map k a) -> Map k a
1800 unions ts
1801 = Foldable.foldl' union empty ts
1802 #if __GLASGOW_HASKELL__
1803 {-# INLINABLE unions #-}
1804 #endif
1805
1806 -- | The union of a list of maps, with a combining operation:
1807 -- (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@).
1808 --
1809 -- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
1810 -- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")]
1811
1812 unionsWith :: (Foldable f, Ord k) => (a->a->a) -> f (Map k a) -> Map k a
1813 unionsWith f ts
1814 = Foldable.foldl' (unionWith f) empty ts
1815 #if __GLASGOW_HASKELL__
1816 {-# INLINABLE unionsWith #-}
1817 #endif
1818
1819 -- | /O(m*log(n\/m + 1)), m <= n/.
1820 -- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and @t2@.
1821 -- It prefers @t1@ when duplicate keys are encountered,
1822 -- i.e. (@'union' == 'unionWith' 'const'@).
1823 --
1824 -- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")]
1825
1826 union :: Ord k => Map k a -> Map k a -> Map k a
1827 union t1 Tip = t1
1828 union t1 (Bin _ k x Tip Tip) = insertR k x t1
1829 union (Bin _ k x Tip Tip) t2 = insert k x t2
1830 union Tip t2 = t2
1831 union t1@(Bin _ k1 x1 l1 r1) t2 = case split k1 t2 of
1832 (l2, r2) | l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 -> t1
1833 | otherwise -> link k1 x1 l1l2 r1r2
1834 where !l1l2 = union l1 l2
1835 !r1r2 = union r1 r2
1836 #if __GLASGOW_HASKELL__
1837 {-# INLINABLE union #-}
1838 #endif
1839
1840 {--------------------------------------------------------------------
1841 Union with a combining function
1842 --------------------------------------------------------------------}
1843 -- | /O(m*log(n\/m + 1)), m <= n/. Union with a combining function.
1844 --
1845 -- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")]
1846
1847 unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
1848 -- QuickCheck says pointer equality never happens here.
1849 unionWith _f t1 Tip = t1
1850 unionWith f t1 (Bin _ k x Tip Tip) = insertWithR f k x t1
1851 unionWith f (Bin _ k x Tip Tip) t2 = insertWith f k x t2
1852 unionWith _f Tip t2 = t2
1853 unionWith f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of
1854 (l2, mb, r2) -> case mb of
1855 Nothing -> link k1 x1 l1l2 r1r2
1856 Just x2 -> link k1 (f x1 x2) l1l2 r1r2
1857 where !l1l2 = unionWith f l1 l2
1858 !r1r2 = unionWith f r1 r2
1859 #if __GLASGOW_HASKELL__
1860 {-# INLINABLE unionWith #-}
1861 #endif
1862
1863 -- | /O(m*log(n\/m + 1)), m <= n/.
1864 -- Union with a combining function.
1865 --
1866 -- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value
1867 -- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")]
1868
1869 unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a
1870 unionWithKey _f t1 Tip = t1
1871 unionWithKey f t1 (Bin _ k x Tip Tip) = insertWithKeyR f k x t1
1872 unionWithKey f (Bin _ k x Tip Tip) t2 = insertWithKey f k x t2
1873 unionWithKey _f Tip t2 = t2
1874 unionWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of
1875 (l2, mb, r2) -> case mb of
1876 Nothing -> link k1 x1 l1l2 r1r2
1877 Just x2 -> link k1 (f k1 x1 x2) l1l2 r1r2
1878 where !l1l2 = unionWithKey f l1 l2
1879 !r1r2 = unionWithKey f r1 r2
1880 #if __GLASGOW_HASKELL__
1881 {-# INLINABLE unionWithKey #-}
1882 #endif
1883
1884 {--------------------------------------------------------------------
1885 Difference
1886 --------------------------------------------------------------------}
1887
1888 -- We don't currently attempt to use any pointer equality tricks for
1889 -- 'difference'. To do so, we'd have to match on the first argument
1890 -- and split the second. Unfortunately, the proof of the time bound
1891 -- relies on doing it the way we do, and it's not clear whether that
1892 -- bound holds the other way.
1893
1894 -- | /O(m*log(n\/m + 1)), m <= n/. Difference of two maps.
1895 -- Return elements of the first map not existing in the second map.
1896 --
1897 -- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b"
1898
1899 difference :: Ord k => Map k a -> Map k b -> Map k a
1900 difference Tip _ = Tip
1901 difference t1 Tip = t1
1902 difference t1 (Bin _ k _ l2 r2) = case split k t1 of
1903 (l1, r1)
1904 | size l1l2 + size r1r2 == size t1 -> t1
1905 | otherwise -> link2 l1l2 r1r2
1906 where
1907 !l1l2 = difference l1 l2
1908 !r1r2 = difference r1 r2
1909 #if __GLASGOW_HASKELL__
1910 {-# INLINABLE difference #-}
1911 #endif
1912
1913 -- | /O(m*log(n\/m + 1)), m <= n/. Remove all keys in a 'Set' from a 'Map'.
1914 --
1915 -- @
1916 -- m \`withoutKeys\` s = 'filterWithKey' (\k _ -> k ``Set.notMember`` s) m
1917 -- m \`withoutKeys\` s = m ``difference`` 'fromSet' (const ()) s
1918 -- @
1919 --
1920 -- @since 0.5.8
1921
1922 withoutKeys :: Ord k => Map k a -> Set k -> Map k a
1923 withoutKeys Tip _ = Tip
1924 withoutKeys m Set.Tip = m
1925 withoutKeys m (Set.Bin _ k ls rs) = case splitMember k m of
1926 (lm, b, rm)
1927 | not b && lm' `ptrEq` lm && rm' `ptrEq` rm -> m
1928 | otherwise -> link2 lm' rm'
1929 where
1930 !lm' = withoutKeys lm ls
1931 !rm' = withoutKeys rm rs
1932 #if __GLASGOW_HASKELL__
1933 {-# INLINABLE withoutKeys #-}
1934 #endif
1935
1936 -- | /O(n+m)/. Difference with a combining function.
1937 -- When two equal keys are
1938 -- encountered, the combining function is applied to the values of these keys.
1939 -- If it returns 'Nothing', the element is discarded (proper set difference). If
1940 -- it returns (@'Just' y@), the element is updated with a new value @y@.
1941 --
1942 -- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing
1943 -- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")])
1944 -- > == singleton 3 "b:B"
1945 differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
1946 differenceWith f = merge preserveMissing dropMissing $
1947 zipWithMaybeMatched (\_ x y -> f x y)
1948 #if __GLASGOW_HASKELL__
1949 {-# INLINABLE differenceWith #-}
1950 #endif
1951
1952 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
1953 -- encountered, the combining function is applied to the key and both values.
1954 -- If it returns 'Nothing', the element is discarded (proper set difference). If
1955 -- it returns (@'Just' y@), the element is updated with a new value @y@.
1956 --
1957 -- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing
1958 -- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")])
1959 -- > == singleton 3 "3:b|B"
1960
1961 differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
1962 differenceWithKey f =
1963 merge preserveMissing dropMissing (zipWithMaybeMatched f)
1964 #if __GLASGOW_HASKELL__
1965 {-# INLINABLE differenceWithKey #-}
1966 #endif
1967
1968
1969 {--------------------------------------------------------------------
1970 Intersection
1971 --------------------------------------------------------------------}
1972 -- | /O(m*log(n\/m + 1)), m <= n/. Intersection of two maps.
1973 -- Return data in the first map for the keys existing in both maps.
1974 -- (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@).
1975 --
1976 -- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a"
1977
1978 intersection :: Ord k => Map k a -> Map k b -> Map k a
1979 intersection Tip _ = Tip
1980 intersection _ Tip = Tip
1981 intersection t1@(Bin _ k x l1 r1) t2
1982 | mb = if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1
1983 then t1
1984 else link k x l1l2 r1r2
1985 | otherwise = link2 l1l2 r1r2
1986 where
1987 !(l2, mb, r2) = splitMember k t2
1988 !l1l2 = intersection l1 l2
1989 !r1r2 = intersection r1 r2
1990 #if __GLASGOW_HASKELL__
1991 {-# INLINABLE intersection #-}
1992 #endif
1993
1994 -- | /O(m*log(n\/m + 1)), m <= n/. Restrict a 'Map' to only those keys
1995 -- found in a 'Set'.
1996 --
1997 -- @
1998 -- m \`restrictKeys\` s = 'filterWithKey' (\k _ -> k ``Set.member`` s) m
1999 -- m \`restrictKeys\` s = m ``intersection`` 'fromSet' (const ()) s
2000 -- @
2001 --
2002 -- @since 0.5.8
2003 restrictKeys :: Ord k => Map k a -> Set k -> Map k a
2004 restrictKeys Tip _ = Tip
2005 restrictKeys _ Set.Tip = Tip
2006 restrictKeys m@(Bin _ k x l1 r1) s
2007 | b = if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1
2008 then m
2009 else link k x l1l2 r1r2
2010 | otherwise = link2 l1l2 r1r2
2011 where
2012 !(l2, b, r2) = Set.splitMember k s
2013 !l1l2 = restrictKeys l1 l2
2014 !r1r2 = restrictKeys r1 r2
2015 #if __GLASGOW_HASKELL__
2016 {-# INLINABLE restrictKeys #-}
2017 #endif
2018
2019 -- | /O(m*log(n\/m + 1)), m <= n/. Intersection with a combining function.
2020 --
2021 -- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA"
2022
2023 intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c
2024 -- We have no hope of pointer equality tricks here because every single
2025 -- element in the result will be a thunk.
2026 intersectionWith _f Tip _ = Tip
2027 intersectionWith _f _ Tip = Tip
2028 intersectionWith f (Bin _ k x1 l1 r1) t2 = case mb of
2029 Just x2 -> link k (f x1 x2) l1l2 r1r2
2030 Nothing -> link2 l1l2 r1r2
2031 where
2032 !(l2, mb, r2) = splitLookup k t2
2033 !l1l2 = intersectionWith f l1 l2
2034 !r1r2 = intersectionWith f r1 r2
2035 #if __GLASGOW_HASKELL__
2036 {-# INLINABLE intersectionWith #-}
2037 #endif
2038
2039 -- | /O(m*log(n\/m + 1)), m <= n/. Intersection with a combining function.
2040 --
2041 -- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar
2042 -- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A"
2043
2044 intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c
2045 intersectionWithKey _f Tip _ = Tip
2046 intersectionWithKey _f _ Tip = Tip
2047 intersectionWithKey f (Bin _ k x1 l1 r1) t2 = case mb of
2048 Just x2 -> link k (f k x1 x2) l1l2 r1r2
2049 Nothing -> link2 l1l2 r1r2
2050 where
2051 !(l2, mb, r2) = splitLookup k t2
2052 !l1l2 = intersectionWithKey f l1 l2
2053 !r1r2 = intersectionWithKey f r1 r2
2054 #if __GLASGOW_HASKELL__
2055 {-# INLINABLE intersectionWithKey #-}
2056 #endif
2057
2058 #if !MIN_VERSION_base (4,8,0)
2059 -- | The identity type.
2060 newtype Identity a = Identity { runIdentity :: a }
2061 #if __GLASGOW_HASKELL__ == 708
2062 instance Functor Identity where
2063 fmap = coerce
2064 instance Applicative Identity where
2065 (<*>) = coerce
2066 pure = Identity
2067 #else
2068 instance Functor Identity where
2069 fmap f (Identity a) = Identity (f a)
2070 instance Applicative Identity where
2071 Identity f <*> Identity x = Identity (f x)
2072 pure = Identity
2073 #endif
2074 #endif
2075
2076 -- | A tactic for dealing with keys present in one map but not the other in
2077 -- 'merge' or 'mergeA'.
2078 --
2079 -- A tactic of type @ WhenMissing f k x z @ is an abstract representation
2080 -- of a function of type @ k -> x -> f (Maybe z) @.
2081 --
2082 -- @since 0.5.9
2083
2084 data WhenMissing f k x y = WhenMissing
2085 { missingSubtree :: Map k x -> f (Map k y)
2086 , missingKey :: k -> x -> f (Maybe y)}
2087
2088 -- | @since 0.5.9
2089 instance (Applicative f, Monad f) => Functor (WhenMissing f k x) where
2090 fmap = mapWhenMissing
2091 {-# INLINE fmap #-}
2092
2093 -- | @since 0.5.9
2094 instance (Applicative f, Monad f)
2095 => Category.Category (WhenMissing f k) where
2096 id = preserveMissing
2097 f . g = traverseMaybeMissing $
2098 \ k x -> missingKey g k x >>= \y ->
2099 case y of
2100 Nothing -> pure Nothing
2101 Just q -> missingKey f k q
2102 {-# INLINE id #-}
2103 {-# INLINE (.) #-}
2104
2105 -- | Equivalent to @ ReaderT k (ReaderT x (MaybeT f)) @.
2106 --
2107 -- @since 0.5.9
2108 instance (Applicative f, Monad f) => Applicative (WhenMissing f k x) where
2109 pure x = mapMissing (\ _ _ -> x)
2110 f <*> g = traverseMaybeMissing $ \k x -> do
2111 res1 <- missingKey f k x
2112 case res1 of
2113 Nothing -> pure Nothing
2114 Just r -> (pure $!) . fmap r =<< missingKey g k x
2115 {-# INLINE pure #-}
2116 {-# INLINE (<*>) #-}
2117
2118 -- | Equivalent to @ ReaderT k (ReaderT x (MaybeT f)) @.
2119 --
2120 -- @since 0.5.9
2121 instance (Applicative f, Monad f) => Monad (WhenMissing f k x) where
2122 #if !MIN_VERSION_base(4,8,0)
2123 return = pure
2124 #endif
2125 m >>= f = traverseMaybeMissing $ \k x -> do
2126 res1 <- missingKey m k x
2127 case res1 of
2128 Nothing -> pure Nothing
2129 Just r -> missingKey (f r) k x
2130 {-# INLINE (>>=) #-}
2131
2132 -- | Map covariantly over a @'WhenMissing' f k x@.
2133 --
2134 -- @since 0.5.9
2135 mapWhenMissing :: (Applicative f, Monad f)
2136 => (a -> b)
2137 -> WhenMissing f k x a -> WhenMissing f k x b
2138 mapWhenMissing f t = WhenMissing
2139 { missingSubtree = \m -> missingSubtree t m >>= \m' -> pure $! fmap f m'
2140 , missingKey = \k x -> missingKey t k x >>= \q -> (pure $! fmap f q) }
2141 {-# INLINE mapWhenMissing #-}
2142
2143 -- | Map covariantly over a @'WhenMissing' f k x@, using only a 'Functor f'
2144 -- constraint.
2145 mapGentlyWhenMissing :: Functor f
2146 => (a -> b)
2147 -> WhenMissing f k x a -> WhenMissing f k x b
2148 mapGentlyWhenMissing f t = WhenMissing
2149 { missingSubtree = \m -> fmap f <$> missingSubtree t m
2150 , missingKey = \k x -> fmap f <$> missingKey t k x }
2151 {-# INLINE mapGentlyWhenMissing #-}
2152
2153 -- | Map covariantly over a @'WhenMatched' f k x@, using only a 'Functor f'
2154 -- constraint.
2155 mapGentlyWhenMatched :: Functor f
2156 => (a -> b)
2157 -> WhenMatched f k x y a -> WhenMatched f k x y b
2158 mapGentlyWhenMatched f t = zipWithMaybeAMatched $
2159 \k x y -> fmap f <$> runWhenMatched t k x y
2160 {-# INLINE mapGentlyWhenMatched #-}
2161
2162 -- | Map contravariantly over a @'WhenMissing' f k _ x@.
2163 --
2164 -- @since 0.5.9
2165 lmapWhenMissing :: (b -> a) -> WhenMissing f k a x -> WhenMissing f k b x
2166 lmapWhenMissing f t = WhenMissing
2167 { missingSubtree = \m -> missingSubtree t (fmap f m)
2168 , missingKey = \k x -> missingKey t k (f x) }
2169 {-# INLINE lmapWhenMissing #-}
2170
2171 -- | Map contravariantly over a @'WhenMatched' f k _ y z@.
2172 --
2173 -- @since 0.5.9
2174 contramapFirstWhenMatched :: (b -> a)
2175 -> WhenMatched f k a y z
2176 -> WhenMatched f k b y z
2177 contramapFirstWhenMatched f t = WhenMatched $
2178 \k x y -> runWhenMatched t k (f x) y
2179 {-# INLINE contramapFirstWhenMatched #-}
2180
2181 -- | Map contravariantly over a @'WhenMatched' f k x _ z@.
2182 --
2183 -- @since 0.5.9
2184 contramapSecondWhenMatched :: (b -> a)
2185 -> WhenMatched f k x a z
2186 -> WhenMatched f k x b z
2187 contramapSecondWhenMatched f t = WhenMatched $
2188 \k x y -> runWhenMatched t k x (f y)
2189 {-# INLINE contramapSecondWhenMatched #-}
2190
2191 -- | A tactic for dealing with keys present in one map but not the other in
2192 -- 'merge'.
2193 --
2194 -- A tactic of type @ SimpleWhenMissing k x z @ is an abstract representation
2195 -- of a function of type @ k -> x -> Maybe z @.
2196 --
2197 -- @since 0.5.9
2198 type SimpleWhenMissing = WhenMissing Identity
2199
2200 -- | A tactic for dealing with keys present in both
2201 -- maps in 'merge' or 'mergeA'.
2202 --
2203 -- A tactic of type @ WhenMatched f k x y z @ is an abstract representation
2204 -- of a function of type @ k -> x -> y -> f (Maybe z) @.
2205 --
2206 -- @since 0.5.9
2207 newtype WhenMatched f k x y z = WhenMatched
2208 { matchedKey :: k -> x -> y -> f (Maybe z) }
2209
2210 -- | Along with zipWithMaybeAMatched, witnesses the isomorphism between
2211 -- @WhenMatched f k x y z@ and @k -> x -> y -> f (Maybe z)@.
2212 --
2213 -- @since 0.5.9
2214 runWhenMatched :: WhenMatched f k x y z -> k -> x -> y -> f (Maybe z)
2215 runWhenMatched = matchedKey
2216 {-# INLINE runWhenMatched #-}
2217
2218 -- | Along with traverseMaybeMissing, witnesses the isomorphism between
2219 -- @WhenMissing f k x y@ and @k -> x -> f (Maybe y)@.
2220 --
2221 -- @since 0.5.9
2222 runWhenMissing :: WhenMissing f k x y -> k -> x -> f (Maybe y)
2223 runWhenMissing = missingKey
2224 {-# INLINE runWhenMissing #-}
2225
2226 -- | @since 0.5.9
2227 instance Functor f => Functor (WhenMatched f k x y) where
2228 fmap = mapWhenMatched
2229 {-# INLINE fmap #-}
2230
2231 -- | @since 0.5.9
2232 instance (Monad f, Applicative f) => Category.Category (WhenMatched f k x) where
2233 id = zipWithMatched (\_ _ y -> y)
2234 f . g = zipWithMaybeAMatched $
2235 \k x y -> do
2236 res <- runWhenMatched g k x y
2237 case res of
2238 Nothing -> pure Nothing
2239 Just r -> runWhenMatched f k x r
2240 {-# INLINE id #-}
2241 {-# INLINE (.) #-}
2242
2243 -- | Equivalent to @ ReaderT k (ReaderT x (ReaderT y (MaybeT f))) @
2244 --
2245 -- @since 0.5.9
2246 instance (Monad f, Applicative f) => Applicative (WhenMatched f k x y) where
2247 pure x = zipWithMatched (\_ _ _ -> x)
2248 fs <*> xs = zipWithMaybeAMatched $ \k x y -> do
2249 res <- runWhenMatched fs k x y
2250 case res of
2251 Nothing -> pure Nothing
2252 Just r -> (pure $!) . fmap r =<< runWhenMatched xs k x y
2253 {-# INLINE pure #-}
2254 {-# INLINE (<*>) #-}
2255
2256 -- | Equivalent to @ ReaderT k (ReaderT x (ReaderT y (MaybeT f))) @
2257 --
2258 -- @since 0.5.9
2259 instance (Monad f, Applicative f) => Monad (WhenMatched f k x y) where
2260 #if !MIN_VERSION_base(4,8,0)
2261 return = pure
2262 #endif
2263 m >>= f = zipWithMaybeAMatched $ \k x y -> do
2264 res <- runWhenMatched m k x y
2265 case res of
2266 Nothing -> pure Nothing
2267 Just r -> runWhenMatched (f r) k x y
2268 {-# INLINE (>>=) #-}
2269
2270 -- | Map covariantly over a @'WhenMatched' f k x y@.
2271 --
2272 -- @since 0.5.9
2273 mapWhenMatched :: Functor f
2274 => (a -> b)
2275 -> WhenMatched f k x y a
2276 -> WhenMatched f k x y b
2277 mapWhenMatched f (WhenMatched g) = WhenMatched $ \k x y -> fmap (fmap f) (g k x y)
2278 {-# INLINE mapWhenMatched #-}
2279
2280 -- | A tactic for dealing with keys present in both maps in 'merge'.
2281 --
2282 -- A tactic of type @ SimpleWhenMatched k x y z @ is an abstract representation
2283 -- of a function of type @ k -> x -> y -> Maybe z @.
2284 --
2285 -- @since 0.5.9
2286 type SimpleWhenMatched = WhenMatched Identity
2287
2288 -- | When a key is found in both maps, apply a function to the
2289 -- key and values and use the result in the merged map.
2290 --
2291 -- @
2292 -- zipWithMatched :: (k -> x -> y -> z)
2293 -- -> SimpleWhenMatched k x y z
2294 -- @
2295 --
2296 -- @since 0.5.9
2297 zipWithMatched :: Applicative f
2298 => (k -> x -> y -> z)
2299 -> WhenMatched f k x y z
2300 zipWithMatched f = WhenMatched $ \ k x y -> pure . Just $ f k x y
2301 {-# INLINE zipWithMatched #-}
2302
2303 -- | When a key is found in both maps, apply a function to the
2304 -- key and values to produce an action and use its result in the merged map.
2305 --
2306 -- @since 0.5.9
2307 zipWithAMatched :: Applicative f
2308 => (k -> x -> y -> f z)
2309 -> WhenMatched f k x y z
2310 zipWithAMatched f = WhenMatched $ \ k x y -> Just <$> f k x y
2311 {-# INLINE zipWithAMatched #-}
2312
2313 -- | When a key is found in both maps, apply a function to the
2314 -- key and values and maybe use the result in the merged map.
2315 --
2316 -- @
2317 -- zipWithMaybeMatched :: (k -> x -> y -> Maybe z)
2318 -- -> SimpleWhenMatched k x y z
2319 -- @
2320 --
2321 -- @since 0.5.9
2322 zipWithMaybeMatched :: Applicative f
2323 => (k -> x -> y -> Maybe z)
2324 -> WhenMatched f k x y z
2325 zipWithMaybeMatched f = WhenMatched $ \ k x y -> pure $ f k x y
2326 {-# INLINE zipWithMaybeMatched #-}
2327
2328 -- | When a key is found in both maps, apply a function to the
2329 -- key and values, perform the resulting action, and maybe use
2330 -- the result in the merged map.
2331 --
2332 -- This is the fundamental 'WhenMatched' tactic.
2333 --
2334 -- @since 0.5.9
2335 zipWithMaybeAMatched :: (k -> x -> y -> f (Maybe z))
2336 -> WhenMatched f k x y z
2337 zipWithMaybeAMatched f = WhenMatched $ \ k x y -> f k x y
2338 {-# INLINE zipWithMaybeAMatched #-}
2339
2340 -- | Drop all the entries whose keys are missing from the other
2341 -- map.
2342 --
2343 -- @
2344 -- dropMissing :: SimpleWhenMissing k x y
2345 -- @
2346 --
2347 -- prop> dropMissing = mapMaybeMissing (\_ _ -> Nothing)
2348 --
2349 -- but @dropMissing@ is much faster.
2350 --
2351 -- @since 0.5.9
2352 dropMissing :: Applicative f => WhenMissing f k x y
2353 dropMissing = WhenMissing
2354 { missingSubtree = const (pure Tip)
2355 , missingKey = \_ _ -> pure Nothing }
2356 {-# INLINE dropMissing #-}
2357
2358 -- | Preserve, unchanged, the entries whose keys are missing from
2359 -- the other map.
2360 --
2361 -- @
2362 -- preserveMissing :: SimpleWhenMissing k x x
2363 -- @
2364 --
2365 -- prop> preserveMissing = Merge.Lazy.mapMaybeMissing (\_ x -> Just x)
2366 --
2367 -- but @preserveMissing@ is much faster.
2368 --
2369 -- @since 0.5.9
2370 preserveMissing :: Applicative f => WhenMissing f k x x
2371 preserveMissing = WhenMissing
2372 { missingSubtree = pure
2373 , missingKey = \_ v -> pure (Just v) }
2374 {-# INLINE preserveMissing #-}
2375
2376 -- | Force the entries whose keys are missing from
2377 -- the other map and otherwise preserve them unchanged.
2378 --
2379 -- @
2380 -- preserveMissing' :: SimpleWhenMissing k x x
2381 -- @
2382 --
2383 -- prop> preserveMissing' = Merge.Lazy.mapMaybeMissing (\_ x -> Just $! x)
2384 --
2385 -- but @preserveMissing'@ is quite a bit faster.
2386 --
2387 -- @since 0.5.9
2388 preserveMissing' :: Applicative f => WhenMissing f k x x
2389 preserveMissing' = WhenMissing
2390 { missingSubtree = \t -> pure $! forceTree t `seq` t
2391 , missingKey = \_ v -> pure $! Just $! v }
2392 {-# INLINE preserveMissing' #-}
2393
2394 -- Force all the values in a tree.
2395 forceTree :: Map k a -> ()
2396 forceTree (Bin _ _ v l r) = v `seq` forceTree l `seq` forceTree r `seq` ()
2397 forceTree Tip = ()
2398
2399 -- | Map over the entries whose keys are missing from the other map.
2400 --
2401 -- @
2402 -- mapMissing :: (k -> x -> y) -> SimpleWhenMissing k x y
2403 -- @
2404 --
2405 -- prop> mapMissing f = mapMaybeMissing (\k x -> Just $ f k x)
2406 --
2407 -- but @mapMissing@ is somewhat faster.
2408 --
2409 -- @since 0.5.9
2410 mapMissing :: Applicative f => (k -> x -> y) -> WhenMissing f k x y
2411 mapMissing f = WhenMissing
2412 { missingSubtree = \m -> pure $! mapWithKey f m
2413 , missingKey = \ k x -> pure $ Just (f k x) }
2414 {-# INLINE mapMissing #-}
2415
2416 -- | Map over the entries whose keys are missing from the other map,
2417 -- optionally removing some. This is the most powerful 'SimpleWhenMissing'
2418 -- tactic, but others are usually more efficient.
2419 --
2420 -- @
2421 -- mapMaybeMissing :: (k -> x -> Maybe y) -> SimpleWhenMissing k x y
2422 -- @
2423 --
2424 -- prop> mapMaybeMissing f = traverseMaybeMissing (\k x -> pure (f k x))
2425 --
2426 -- but @mapMaybeMissing@ uses fewer unnecessary 'Applicative' operations.
2427 --
2428 -- @since 0.5.9
2429 mapMaybeMissing :: Applicative f => (k -> x -> Maybe y) -> WhenMissing f k x y
2430 mapMaybeMissing f = WhenMissing
2431 { missingSubtree = \m -> pure $! mapMaybeWithKey f m
2432 , missingKey = \k x -> pure $! f k x }
2433 {-# INLINE mapMaybeMissing #-}
2434
2435 -- | Filter the entries whose keys are missing from the other map.
2436 --
2437 -- @
2438 -- filterMissing :: (k -> x -> Bool) -> SimpleWhenMissing k x x
2439 -- @
2440 --
2441 -- prop> filterMissing f = Merge.Lazy.mapMaybeMissing $ \k x -> guard (f k x) *> Just x
2442 --
2443 -- but this should be a little faster.
2444 --
2445 -- @since 0.5.9
2446 filterMissing :: Applicative f
2447 => (k -> x -> Bool) -> WhenMissing f k x x
2448 filterMissing f = WhenMissing
2449 { missingSubtree = \m -> pure $! filterWithKey f m
2450 , missingKey = \k x -> pure $! if f k x then Just x else Nothing }
2451 {-# INLINE filterMissing #-}
2452
2453 -- | Filter the entries whose keys are missing from the other map
2454 -- using some 'Applicative' action.
2455 --
2456 -- @
2457 -- filterAMissing f = Merge.Lazy.traverseMaybeMissing $
2458 -- \k x -> (\b -> guard b *> Just x) <$> f k x
2459 -- @
2460 --
2461 -- but this should be a little faster.
2462 --
2463 -- @since 0.5.9
2464 filterAMissing :: Applicative f
2465 => (k -> x -> f Bool) -> WhenMissing f k x x
2466 filterAMissing f = WhenMissing
2467 { missingSubtree = \m -> filterWithKeyA f m
2468 , missingKey = \k x -> bool Nothing (Just x) <$> f k x }
2469 {-# INLINE filterAMissing #-}
2470
2471 -- | This wasn't in Data.Bool until 4.7.0, so we define it here
2472 bool :: a -> a -> Bool -> a
2473 bool f _ False = f
2474 bool _ t True = t
2475
2476 -- | Traverse over the entries whose keys are missing from the other map.
2477 --
2478 -- @since 0.5.9
2479 traverseMissing :: Applicative f
2480 => (k -> x -> f y) -> WhenMissing f k x y
2481 traverseMissing f = WhenMissing
2482 { missingSubtree = traverseWithKey f
2483 , missingKey = \k x -> Just <$> f k x }
2484 {-# INLINE traverseMissing #-}
2485
2486 -- | Traverse over the entries whose keys are missing from the other map,
2487 -- optionally producing values to put in the result.
2488 -- This is the most powerful 'WhenMissing' tactic, but others are usually
2489 -- more efficient.
2490 --
2491 -- @since 0.5.9
2492 traverseMaybeMissing :: Applicative f
2493 => (k -> x -> f (Maybe y)) -> WhenMissing f k x y
2494 traverseMaybeMissing f = WhenMissing
2495 { missingSubtree = traverseMaybeWithKey f
2496 , missingKey = f }
2497 {-# INLINE traverseMaybeMissing #-}
2498
2499 -- | Merge two maps.
2500 --
2501 -- 'merge' takes two 'WhenMissing' tactics, a 'WhenMatched'
2502 -- tactic and two maps. It uses the tactics to merge the maps.
2503 -- Its behavior is best understood via its fundamental tactics,
2504 -- 'mapMaybeMissing' and 'zipWithMaybeMatched'.
2505 --
2506 -- Consider
2507 --
2508 -- @
2509 -- merge (mapMaybeMissing g1)
2510 -- (mapMaybeMissing g2)
2511 -- (zipWithMaybeMatched f)
2512 -- m1 m2
2513 -- @
2514 --
2515 -- Take, for example,
2516 --
2517 -- @
2518 -- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')]
2519 -- m2 = [(1, "one"), (2, "two"), (4, "three")]
2520 -- @
2521 --
2522 -- 'merge' will first \"align\" these maps by key:
2523 --
2524 -- @
2525 -- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')]
2526 -- m2 = [(1, "one"), (2, "two"), (4, "three")]
2527 -- @
2528 --
2529 -- It will then pass the individual entries and pairs of entries
2530 -- to @g1@, @g2@, or @f@ as appropriate:
2531 --
2532 -- @
2533 -- maybes = [g1 0 \'a\', f 1 \'b\' "one", g2 2 "two", g1 3 \'c\', f 4 \'d\' "three"]
2534 -- @
2535 --
2536 -- This produces a 'Maybe' for each key:
2537 --
2538 -- @
2539 -- keys = 0 1 2 3 4
2540 -- results = [Nothing, Just True, Just False, Nothing, Just True]
2541 -- @
2542 --
2543 -- Finally, the @Just@ results are collected into a map:
2544 --
2545 -- @
2546 -- return value = [(1, True), (2, False), (4, True)]
2547 -- @
2548 --
2549 -- The other tactics below are optimizations or simplifications of
2550 -- 'mapMaybeMissing' for special cases. Most importantly,
2551 --
2552 -- * 'dropMissing' drops all the keys.
2553 -- * 'preserveMissing' leaves all the entries alone.
2554 --
2555 -- When 'merge' is given three arguments, it is inlined at the call
2556 -- site. To prevent excessive inlining, you should typically use 'merge'
2557 -- to define your custom combining functions.
2558 --
2559 --
2560 -- Examples:
2561 --
2562 -- prop> unionWithKey f = merge preserveMissing preserveMissing (zipWithMatched f)
2563 -- prop> intersectionWithKey f = merge dropMissing dropMissing (zipWithMatched f)
2564 -- prop> differenceWith f = merge diffPreserve diffDrop f
2565 -- prop> symmetricDifference = merge diffPreserve diffPreserve (\ _ _ _ -> Nothing)
2566 -- prop> mapEachPiece f g h = merge (diffMapWithKey f) (diffMapWithKey g)
2567 --
2568 -- @since 0.5.9
2569 merge :: Ord k
2570 => SimpleWhenMissing k a c -- ^ What to do with keys in @m1@ but not @m2@
2571 -> SimpleWhenMissing k b c -- ^ What to do with keys in @m2@ but not @m1@
2572 -> SimpleWhenMatched k a b c -- ^ What to do with keys in both @m1@ and @m2@
2573 -> Map k a -- ^ Map @m1@
2574 -> Map k b -- ^ Map @m2@
2575 -> Map k c
2576 merge g1 g2 f m1 m2 = runIdentity $
2577 mergeA g1 g2 f m1 m2
2578 {-# INLINE merge #-}
2579
2580 -- | An applicative version of 'merge'.
2581 --
2582 -- 'mergeA' takes two 'WhenMissing' tactics, a 'WhenMatched'
2583 -- tactic and two maps. It uses the tactics to merge the maps.
2584 -- Its behavior is best understood via its fundamental tactics,
2585 -- 'traverseMaybeMissing' and 'zipWithMaybeAMatched'.
2586 --
2587 -- Consider
2588 --
2589 -- @
2590 -- mergeA (traverseMaybeMissing g1)
2591 -- (traverseMaybeMissing g2)
2592 -- (zipWithMaybeAMatched f)
2593 -- m1 m2
2594 -- @
2595 --
2596 -- Take, for example,
2597 --
2598 -- @
2599 -- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')]
2600 -- m2 = [(1, "one"), (2, "two"), (4, "three")]
2601 -- @
2602 --
2603 -- @mergeA@ will first \"align\" these maps by key:
2604 --
2605 -- @
2606 -- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')]
2607 -- m2 = [(1, "one"), (2, "two"), (4, "three")]
2608 -- @
2609 --
2610 -- It will then pass the individual entries and pairs of entries
2611 -- to @g1@, @g2@, or @f@ as appropriate:
2612 --
2613 -- @
2614 -- actions = [g1 0 \'a\', f 1 \'b\' "one", g2 2 "two", g1 3 \'c\', f 4 \'d\' "three"]
2615 -- @
2616 --
2617 -- Next, it will perform the actions in the @actions@ list in order from
2618 -- left to right.
2619 --
2620 -- @
2621 -- keys = 0 1 2 3 4
2622 -- results = [Nothing, Just True, Just False, Nothing, Just True]
2623 -- @
2624 --
2625 -- Finally, the @Just@ results are collected into a map:
2626 --
2627 -- @
2628 -- return value = [(1, True), (2, False), (4, True)]
2629 -- @
2630 --
2631 -- The other tactics below are optimizations or simplifications of
2632 -- 'traverseMaybeMissing' for special cases. Most importantly,
2633 --
2634 -- * 'dropMissing' drops all the keys.
2635 -- * 'preserveMissing' leaves all the entries alone.
2636 -- * 'mapMaybeMissing' does not use the 'Applicative' context.
2637 --
2638 -- When 'mergeA' is given three arguments, it is inlined at the call
2639 -- site. To prevent excessive inlining, you should generally only use
2640 -- 'mergeA' to define custom combining functions.
2641 --
2642 -- @since 0.5.9
2643 mergeA
2644 :: (Applicative f, Ord k)
2645 => WhenMissing f k a c -- ^ What to do with keys in @m1@ but not @m2@
2646 -> WhenMissing f k b c -- ^ What to do with keys in @m2@ but not @m1@
2647 -> WhenMatched f k a b c -- ^ What to do with keys in both @m1@ and @m2@
2648 -> Map k a -- ^ Map @m1@
2649 -> Map k b -- ^ Map @m2@
2650 -> f (Map k c)
2651 mergeA
2652 WhenMissing{missingSubtree = g1t, missingKey = g1k}
2653 WhenMissing{missingSubtree = g2t}
2654 (WhenMatched f) = go
2655 where
2656 go t1 Tip = g1t t1
2657 go Tip t2 = g2t t2
2658 go (Bin _ kx x1 l1 r1) t2 = case splitLookup kx t2 of
2659 (l2, mx2, r2) -> case mx2 of
2660 Nothing -> liftA3 (\l' mx' r' -> maybe link2 (link kx) mx' l' r')
2661 l1l2 (g1k kx x1) r1r2
2662 Just x2 -> liftA3 (\l' mx' r' -> maybe link2 (link kx) mx' l' r')
2663 l1l2 (f kx x1 x2) r1r2
2664 where
2665 !l1l2 = go l1 l2
2666 !r1r2 = go r1 r2
2667 {-# INLINE mergeA #-}
2668
2669
2670 {--------------------------------------------------------------------
2671 MergeWithKey
2672 --------------------------------------------------------------------}
2673
2674 -- | /O(n+m)/. An unsafe general combining function.
2675 --
2676 -- WARNING: This function can produce corrupt maps and its results
2677 -- may depend on the internal structures of its inputs. Users should
2678 -- prefer 'merge' or 'mergeA'.
2679 --
2680 -- When 'mergeWithKey' is given three arguments, it is inlined to the call
2681 -- site. You should therefore use 'mergeWithKey' only to define custom
2682 -- combining functions. For example, you could define 'unionWithKey',
2683 -- 'differenceWithKey' and 'intersectionWithKey' as
2684 --
2685 -- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2
2686 -- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2
2687 -- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2
2688 --
2689 -- When calling @'mergeWithKey' combine only1 only2@, a function combining two
2690 -- 'Map's is created, such that
2691 --
2692 -- * if a key is present in both maps, it is passed with both corresponding
2693 -- values to the @combine@ function. Depending on the result, the key is either
2694 -- present in the result with specified value, or is left out;
2695 --
2696 -- * a nonempty subtree present only in the first map is passed to @only1@ and
2697 -- the output is added to the result;
2698 --
2699 -- * a nonempty subtree present only in the second map is passed to @only2@ and
2700 -- the output is added to the result.
2701 --
2702 -- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/.
2703 -- The values can be modified arbitrarily. Most common variants of @only1@ and
2704 -- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@,
2705 -- @'filterWithKey' f@, or @'mapMaybeWithKey' f@ could be used for any @f@.
2706
2707 mergeWithKey :: Ord k
2708 => (k -> a -> b -> Maybe c)
2709 -> (Map k a -> Map k c)
2710 -> (Map k b -> Map k c)
2711 -> Map k a -> Map k b -> Map k c
2712 mergeWithKey f g1 g2 = go
2713 where
2714 go Tip t2 = g2 t2
2715 go t1 Tip = g1 t1
2716 go (Bin _ kx x l1 r1) t2 =
2717 case found of
2718 Nothing -> case g1 (singleton kx x) of
2719 Tip -> link2 l' r'
2720 (Bin _ _ x' Tip Tip) -> link kx x' l' r'
2721 _ -> error "mergeWithKey: Given function only1 does not fulfill required conditions (see documentation)"
2722 Just x2 -> case f kx x x2 of
2723 Nothing -> link2 l' r'
2724 Just x' -> link kx x' l' r'
2725 where
2726 (l2, found, r2) = splitLookup kx t2
2727 l' = go l1 l2
2728 r' = go r1 r2
2729 {-# INLINE mergeWithKey #-}
2730
2731 {--------------------------------------------------------------------
2732 Submap
2733 --------------------------------------------------------------------}
2734 -- | /O(m*log(n\/m + 1)), m <= n/.
2735 -- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
2736 --
2737 isSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool
2738 isSubmapOf m1 m2 = isSubmapOfBy (==) m1 m2
2739 #if __GLASGOW_HASKELL__
2740 {-# INLINABLE isSubmapOf #-}
2741 #endif
2742
2743 {- | /O(m*log(n\/m + 1)), m <= n/.
2744 The expression (@'isSubmapOfBy' f t1 t2@) returns 'True' if
2745 all keys in @t1@ are in tree @t2@, and when @f@ returns 'True' when
2746 applied to their respective values. For example, the following
2747 expressions are all 'True':
2748
2749 > isSubmapOfBy (==) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
2750 > isSubmapOfBy (<=) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
2751 > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1),('b',2)])
2752
2753 But the following are all 'False':
2754
2755 > isSubmapOfBy (==) (fromList [('a',2)]) (fromList [('a',1),('b',2)])
2756 > isSubmapOfBy (<) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
2757 > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1)])
2758
2759 Note that @isSubmapOfBy (\_ _ -> True) m1 m2@ tests whether all the keys
2760 in @m1@ are also keys in @m2@.
2761
2762 -}
2763 isSubmapOfBy :: Ord k => (a->b->Bool) -> Map k a -> Map k b -> Bool
2764 isSubmapOfBy f t1 t2
2765 = size t1 <= size t2 && submap' f t1 t2
2766 #if __GLASGOW_HASKELL__
2767 {-# INLINABLE isSubmapOfBy #-}
2768 #endif
2769
2770 -- Test whether a map is a submap of another without the *initial*
2771 -- size test. See Data.Set.Internal.isSubsetOfX for notes on
2772 -- implementation and analysis.
2773 submap' :: Ord a => (b -> c -> Bool) -> Map a b -> Map a c -> Bool
2774 submap' _ Tip _ = True
2775 submap' _ _ Tip = False
2776 submap' f (Bin 1 kx x _ _) t
2777 = case lookup kx t of
2778 Just y -> f x y
2779 Nothing -> False
2780 submap' f (Bin _ kx x l r) t
2781 = case found of
2782 Nothing -> False
2783 Just y -> f x y
2784 && size l <= size lt && size r <= size gt
2785 && submap' f l lt && submap' f r gt
2786 where
2787 (lt,found,gt) = splitLookup kx t
2788 #if __GLASGOW_HASKELL__
2789 {-# INLINABLE submap' #-}
2790 #endif
2791
2792 -- | /O(m*log(n\/m + 1)), m <= n/. Is this a proper submap? (ie. a submap but not equal).
2793 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
2794 isProperSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool
2795 isProperSubmapOf m1 m2
2796 = isProperSubmapOfBy (==) m1 m2
2797 #if __GLASGOW_HASKELL__
2798 {-# INLINABLE isProperSubmapOf #-}
2799 #endif
2800
2801 {- | /O(m*log(n\/m + 1)), m <= n/. Is this a proper submap? (ie. a submap but not equal).
2802 The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
2803 @m1@ and @m2@ are not equal,
2804 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
2805 applied to their respective values. For example, the following
2806 expressions are all 'True':
2807
2808 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
2809 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
2810
2811 But the following are all 'False':
2812
2813 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
2814 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
2815 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
2816
2817
2818 -}
2819 isProperSubmapOfBy :: Ord k => (a -> b -> Bool) -> Map k a -> Map k b -> Bool
2820 isProperSubmapOfBy f t1 t2
2821 = size t1 < size t2 && submap' f t1 t2
2822 #if __GLASGOW_HASKELL__
2823 {-# INLINABLE isProperSubmapOfBy #-}
2824 #endif
2825
2826 {--------------------------------------------------------------------
2827 Filter and partition
2828 --------------------------------------------------------------------}
2829 -- | /O(n)/. Filter all values that satisfy the predicate.
2830 --
2831 -- > filter (> "a") (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
2832 -- > filter (> "x") (fromList [(5,"a"), (3,"b")]) == empty
2833 -- > filter (< "a") (fromList [(5,"a"), (3,"b")]) == empty
2834
2835 filter :: (a -> Bool) -> Map k a -> Map k a
2836 filter p m
2837 = filterWithKey (\_ x -> p x) m
2838
2839 -- | /O(n)/. Filter all keys\/values that satisfy the predicate.
2840 --
2841 -- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
2842
2843 filterWithKey :: (k -> a -> Bool) -> Map k a -> Map k a
2844 filterWithKey _ Tip = Tip
2845 filterWithKey p t@(Bin _ kx x l r)
2846 | p kx x = if pl `ptrEq` l && pr `ptrEq` r
2847 then t
2848 else link kx x pl pr
2849 | otherwise = link2 pl pr
2850 where !pl = filterWithKey p l
2851 !pr = filterWithKey p r
2852
2853 -- | /O(n)/. Filter keys and values using an 'Applicative'
2854 -- predicate.
2855 filterWithKeyA :: Applicative f => (k -> a -> f Bool) -> Map k a -> f (Map k a)
2856 filterWithKeyA _ Tip = pure Tip
2857 filterWithKeyA p t@(Bin _ kx x l r) =
2858 liftA3 combine (p kx x) (filterWithKeyA p l) (filterWithKeyA p r)
2859 where
2860 combine True pl pr
2861 | pl `ptrEq` l && pr `ptrEq` r = t
2862 | otherwise = link kx x pl pr
2863 combine False pl pr = link2 pl pr
2864
2865 -- | /O(log n)/. Take while a predicate on the keys holds.
2866 -- The user is responsible for ensuring that for all keys @j@ and @k@ in the map,
2867 -- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'.
2868 --
2869 -- @
2870 -- takeWhileAntitone p = 'fromDistinctAscList' . 'Data.List.takeWhile' (p . fst) . 'toList'
2871 -- takeWhileAntitone p = 'filterWithKey' (\k _ -> p k)
2872 -- @
2873 --
2874 -- @since 0.5.8
2875
2876 takeWhileAntitone :: (k -> Bool) -> Map k a -> Map k a
2877 takeWhileAntitone _ Tip = Tip
2878 takeWhileAntitone p (Bin _ kx x l r)
2879 | p kx = link kx x l (takeWhileAntitone p r)
2880 | otherwise = takeWhileAntitone p l
2881
2882 -- | /O(log n)/. Drop while a predicate on the keys holds.
2883 -- The user is responsible for ensuring that for all keys @j@ and @k@ in the map,
2884 -- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'.
2885 --
2886 -- @
2887 -- dropWhileAntitone p = 'fromDistinctAscList' . 'Data.List.dropWhile' (p . fst) . 'toList'
2888 -- dropWhileAntitone p = 'filterWithKey' (\k -> not (p k))
2889 -- @
2890 --
2891 -- @since 0.5.8
2892
2893 dropWhileAntitone :: (k -> Bool) -> Map k a -> Map k a
2894 dropWhileAntitone _ Tip = Tip
2895 dropWhileAntitone p (Bin _ kx x l r)
2896 | p kx = dropWhileAntitone p r
2897 | otherwise = link kx x (dropWhileAntitone p l) r
2898
2899 -- | /O(log n)/. Divide a map at the point where a predicate on the keys stops holding.
2900 -- The user is responsible for ensuring that for all keys @j@ and @k@ in the map,
2901 -- @j \< k ==\> p j \>= p k@.
2902 --
2903 -- @
2904 -- spanAntitone p xs = ('takeWhileAntitone' p xs, 'dropWhileAntitone' p xs)
2905 -- spanAntitone p xs = partitionWithKey (\k _ -> p k) xs
2906 -- @
2907 --
2908 -- Note: if @p@ is not actually antitone, then @spanAntitone@ will split the map
2909 -- at some /unspecified/ point where the predicate switches from holding to not
2910 -- holding (where the predicate is seen to hold before the first key and to fail
2911 -- after the last key).
2912 --
2913 -- @since 0.5.8
2914
2915 spanAntitone :: (k -> Bool) -> Map k a -> (Map k a, Map k a)
2916 spanAntitone p0 m = toPair (go p0 m)
2917 where
2918 go _ Tip = Tip :*: Tip
2919 go p (Bin _ kx x l r)
2920 | p kx = let u :*: v = go p r in link kx x l u :*: v
2921 | otherwise = let u :*: v = go p l in u :*: link kx x v r
2922
2923 -- | /O(n)/. Partition the map according to a predicate. The first
2924 -- map contains all elements that satisfy the predicate, the second all
2925 -- elements that fail the predicate. See also 'split'.
2926 --
2927 -- > partition (> "a") (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
2928 -- > partition (< "x") (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
2929 -- > partition (> "x") (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
2930
2931 partition :: (a -> Bool) -> Map k a -> (Map k a,Map k a)
2932 partition p m
2933 = partitionWithKey (\_ x -> p x) m
2934
2935 -- | /O(n)/. Partition the map according to a predicate. The first
2936 -- map contains all elements that satisfy the predicate, the second all
2937 -- elements that fail the predicate. See also 'split'.
2938 --
2939 -- > partitionWithKey (\ k _ -> k > 3) (fromList [(5,"a"), (3,"b")]) == (singleton 5 "a", singleton 3 "b")
2940 -- > partitionWithKey (\ k _ -> k < 7) (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
2941 -- > partitionWithKey (\ k _ -> k > 7) (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
2942
2943 partitionWithKey :: (k -> a -> Bool) -> Map k a -> (Map k a,Map k a)
2944 partitionWithKey p0 t0 = toPair $ go p0 t0
2945 where
2946 go _ Tip = (Tip :*: Tip)
2947 go p t@(Bin _ kx x l r)
2948 | p kx x = (if l1 `ptrEq` l && r1 `ptrEq` r
2949 then t
2950 else link kx x l1 r1) :*: link2 l2 r2
2951 | otherwise = link2 l1 r1 :*:
2952 (if l2 `ptrEq` l && r2 `ptrEq` r
2953 then t
2954 else link kx x l2 r2)
2955 where
2956 (l1 :*: l2) = go p l
2957 (r1 :*: r2) = go p r
2958
2959 -- | /O(n)/. Map values and collect the 'Just' results.
2960 --
2961 -- > let f x = if x == "a" then Just "new a" else Nothing
2962 -- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a"
2963
2964 mapMaybe :: (a -> Maybe b) -> Map k a -> Map k b
2965 mapMaybe f = mapMaybeWithKey (\_ x -> f x)
2966
2967 -- | /O(n)/. Map keys\/values and collect the 'Just' results.
2968 --
2969 -- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing
2970 -- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3"
2971
2972 mapMaybeWithKey :: (k -> a -> Maybe b) -> Map k a -> Map k b
2973 mapMaybeWithKey _ Tip = Tip
2974 mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of
2975 Just y -> link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r)
2976 Nothing -> link2 (mapMaybeWithKey f l) (mapMaybeWithKey f r)
2977
2978 -- | /O(n)/. Traverse keys\/values and collect the 'Just' results.
2979 --
2980 -- @since 0.5.8
2981 traverseMaybeWithKey :: Applicative f
2982 => (k -> a -> f (Maybe b)) -> Map k a -> f (Map k b)
2983 traverseMaybeWithKey = go
2984 where
2985 go _ Tip = pure Tip
2986 go f (Bin _ kx x Tip Tip) = maybe Tip (\x' -> Bin 1 kx x' Tip Tip) <$> f kx x
2987 go f (Bin _ kx x l r) = liftA3 combine (go f l) (f kx x) (go f r)
2988 where
2989 combine !l' mx !r' = case mx of
2990 Nothing -> link2 l' r'
2991 Just x' -> link kx x' l' r'
2992
2993 -- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
2994 --
2995 -- > let f a = if a < "c" then Left a else Right a
2996 -- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
2997 -- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")])
2998 -- >
2999 -- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
3000 -- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
3001
3002 mapEither :: (a -> Either b c) -> Map k a -> (Map k b, Map k c)
3003 mapEither f m
3004 = mapEitherWithKey (\_ x -> f x) m
3005
3006 -- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
3007 --
3008 -- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a)
3009 -- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
3010 -- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")])
3011 -- >
3012 -- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
3013 -- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")])
3014
3015 mapEitherWithKey :: (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c)
3016 mapEitherWithKey f0 t0 = toPair $ go f0 t0
3017 where
3018 go _ Tip = (Tip :*: Tip)
3019 go f (Bin _ kx x l r) = case f kx x of
3020 Left y -> link kx y l1 r1 :*: link2 l2 r2
3021 Right z -> link2 l1 r1 :*: link kx z l2 r2
3022 where
3023 (l1 :*: l2) = go f l
3024 (r1 :*: r2) = go f r
3025
3026 {--------------------------------------------------------------------
3027 Mapping
3028 --------------------------------------------------------------------}
3029 -- | /O(n)/. Map a function over all values in the map.
3030 --
3031 -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")]
3032
3033 map :: (a -> b) -> Map k a -> Map k b
3034 map f = go where
3035 go Tip = Tip
3036 go (Bin sx kx x l r) = Bin sx kx (f x) (go l) (go r)
3037 -- We use a `go` function to allow `map` to inline. This makes
3038 -- a big difference if someone uses `map (const x) m` instead
3039 -- of `x <$ m`; it doesn't seem to do any harm.
3040
3041 #ifdef __GLASGOW_HASKELL__
3042 {-# NOINLINE [1] map #-}
3043 {-# RULES
3044 "map/map" forall f g xs . map f (map g xs) = map (f . g) xs
3045 #-}
3046 #endif
3047 #if __GLASGOW_HASKELL__ >= 709
3048 -- Safe coercions were introduced in 7.8, but did not work well with RULES yet.
3049 {-# RULES
3050 "map/coerce" map coerce = coerce
3051 #-}
3052 #endif
3053
3054 -- | /O(n)/. Map a function over all values in the map.
3055 --
3056 -- > let f key x = (show key) ++ ":" ++ x
3057 -- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")]
3058
3059 mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
3060 mapWithKey _ Tip = Tip
3061 mapWithKey f (Bin sx kx x l r) = Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey f r)
3062
3063 #ifdef __GLASGOW_HASKELL__
3064 {-# NOINLINE [1] mapWithKey #-}
3065 {-# RULES
3066 "mapWithKey/mapWithKey" forall f g xs . mapWithKey f (mapWithKey g xs) =
3067 mapWithKey (\k a -> f k (g k a)) xs
3068 "mapWithKey/map" forall f g xs . mapWithKey f (map g xs) =
3069 mapWithKey (\k a -> f k (g a)) xs
3070 "map/mapWithKey" forall f g xs . map f (mapWithKey g xs) =
3071 mapWithKey (\k a -> f (g k a)) xs
3072 #-}
3073 #endif
3074
3075 -- | /O(n)/.
3076 -- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@
3077 -- That is, behaves exactly like a regular 'traverse' except that the traversing
3078 -- function also has access to the key associated with a value.
3079 --
3080 -- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(1, 'a'), (5, 'e')]) == Just (fromList [(1, 'b'), (5, 'f')])
3081 -- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(2, 'c')]) == Nothing
3082 traverseWithKey :: Applicative t => (k -> a -> t b) -> Map k a -> t (Map k b)
3083 traverseWithKey f = go
3084 where
3085 go Tip = pure Tip
3086 go (Bin 1 k v _ _) = (\v' -> Bin 1 k v' Tip Tip) <$> f k v
3087 go (Bin s k v l r) = liftA3 (flip (Bin s k)) (go l) (f k v) (go r)
3088 {-# INLINE traverseWithKey #-}
3089
3090 -- | /O(n)/. The function 'mapAccum' threads an accumulating
3091 -- argument through the map in ascending order of keys.
3092 --
3093 -- > let f a b = (a ++ b, b ++ "X")
3094 -- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")])
3095
3096 mapAccum :: (a -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
3097 mapAccum f a m
3098 = mapAccumWithKey (\a' _ x' -> f a' x') a m
3099
3100 -- | /O(n)/. The function 'mapAccumWithKey' threads an accumulating
3101 -- argument through the map in ascending order of keys.
3102 --
3103 -- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X")
3104 -- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")])
3105
3106 mapAccumWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
3107 mapAccumWithKey f a t
3108 = mapAccumL f a t
3109
3110 -- | /O(n)/. The function 'mapAccumL' threads an accumulating
3111 -- argument through the map in ascending order of keys.
3112 mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
3113 mapAccumL _ a Tip = (a,Tip)
3114 mapAccumL f a (Bin sx kx x l r) =
3115 let (a1,l') = mapAccumL f a l
3116 (a2,x') = f a1 kx x
3117 (a3,r') = mapAccumL f a2 r
3118 in (a3,Bin sx kx x' l' r')
3119
3120 -- | /O(n)/. The function 'mapAccumR' threads an accumulating
3121 -- argument through the map in descending order of keys.
3122 mapAccumRWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
3123 mapAccumRWithKey _ a Tip = (a,Tip)
3124 mapAccumRWithKey f a (Bin sx kx x l r) =
3125 let (a1,r') = mapAccumRWithKey f a r
3126 (a2,x') = f a1 kx x
3127 (a3,l') = mapAccumRWithKey f a2 l
3128 in (a3,Bin sx kx x' l' r')
3129
3130 -- | /O(n*log n)/.
3131 -- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@.
3132 --
3133 -- The size of the result may be smaller if @f@ maps two or more distinct
3134 -- keys to the same new key. In this case the value at the greatest of the
3135 -- original keys is retained.
3136 --
3137 -- > mapKeys (+ 1) (fromList [(5,"a"), (3,"b")]) == fromList [(4, "b"), (6, "a")]
3138 -- > mapKeys (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "c"
3139 -- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c"
3140
3141 mapKeys :: Ord k2 => (k1->k2) -> Map k1 a -> Map k2 a
3142 mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) []
3143 #if __GLASGOW_HASKELL__
3144 {-# INLINABLE mapKeys #-}
3145 #endif
3146
3147 -- | /O(n*log n)/.
3148 -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@.
3149 --
3150 -- The size of the result may be smaller if @f@ maps two or more distinct
3151 -- keys to the same new key. In this case the associated values will be
3152 -- combined using @c@. The value at the greater of the two original keys
3153 -- is used as the first argument to @c@.
3154 --
3155 -- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab"
3156 -- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab"
3157
3158 mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a
3159 mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) []
3160 #if __GLASGOW_HASKELL__
3161 {-# INLINABLE mapKeysWith #-}
3162 #endif
3163
3164
3165 -- | /O(n)/.
3166 -- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@
3167 -- is strictly monotonic.
3168 -- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@.
3169 -- /The precondition is not checked./
3170 -- Semi-formally, we have:
3171 --
3172 -- > and [x < y ==> f x < f y | x <- ls, y <- ls]
3173 -- > ==> mapKeysMonotonic f s == mapKeys f s
3174 -- > where ls = keys s
3175 --
3176 -- This means that @f@ maps distinct original keys to distinct resulting keys.
3177 -- This function has better performance than 'mapKeys'.
3178 --
3179 -- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")]
3180 -- > valid (mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")])) == True
3181 -- > valid (mapKeysMonotonic (\ _ -> 1) (fromList [(5,"a"), (3,"b")])) == False
3182
3183 mapKeysMonotonic :: (k1->k2) -> Map k1 a -> Map k2 a
3184 mapKeysMonotonic _ Tip = Tip
3185 mapKeysMonotonic f (Bin sz k x l r) =
3186 Bin sz (f k) x (mapKeysMonotonic f l) (mapKeysMonotonic f r)
3187
3188 {--------------------------------------------------------------------
3189 Folds
3190 --------------------------------------------------------------------}
3191
3192 -- | /O(n)/. Fold the values in the map using the given right-associative
3193 -- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'elems'@.
3194 --
3195 -- For example,
3196 --
3197 -- > elems map = foldr (:) [] map
3198 --
3199 -- > let f a len = len + (length a)
3200 -- > foldr f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
3201 foldr :: (a -> b -> b) -> b -> Map k a -> b
3202 foldr f z = go z
3203 where
3204 go z' Tip = z'
3205 go z' (Bin _ _ x l r) = go (f x (go z' r)) l
3206 {-# INLINE foldr #-}
3207
3208 -- | /O(n)/. A strict version of 'foldr'. Each application of the operator is
3209 -- evaluated before using the result in the next application. This
3210 -- function is strict in the starting value.
3211 foldr' :: (a -> b -> b) -> b -> Map k a -> b
3212 foldr' f z = go z
3213 where
3214 go !z' Tip = z'
3215 go z' (Bin _ _ x l r) = go (f x (go z' r)) l
3216 {-# INLINE foldr' #-}
3217
3218 -- | /O(n)/. Fold the values in the map using the given left-associative
3219 -- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@.
3220 --
3221 -- For example,
3222 --
3223 -- > elems = reverse . foldl (flip (:)) []
3224 --
3225 -- > let f len a = len + (length a)
3226 -- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
3227 foldl :: (a -> b -> a) -> a -> Map k b -> a
3228 foldl f z = go z
3229 where
3230 go z' Tip = z'
3231 go z' (Bin _ _ x l r) = go (f (go z' l) x) r
3232 {-# INLINE foldl #-}
3233
3234 -- | /O(n)/. A strict version of 'foldl'. Each application of the operator is
3235 -- evaluated before using the result in the next application. This
3236 -- function is strict in the starting value.
3237 foldl' :: (a -> b -> a) -> a -> Map k b -> a
3238 foldl' f z = go z
3239 where
3240 go !z' Tip = z'
3241 go z' (Bin _ _ x l r) = go (f (go z' l) x) r
3242 {-# INLINE foldl' #-}
3243
3244 -- | /O(n)/. Fold the keys and values in the map using the given right-associative
3245 -- binary operator, such that
3246 -- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
3247 --
3248 -- For example,
3249 --
3250 -- > keys map = foldrWithKey (\k x ks -> k:ks) [] map
3251 --
3252 -- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
3253 -- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)"
3254 foldrWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b
3255 foldrWithKey f z = go z
3256 where
3257 go z' Tip = z'
3258 go z' (Bin _ kx x l r) = go (f kx x (go z' r)) l
3259 {-# INLINE foldrWithKey #-}
3260
3261 -- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is
3262 -- evaluated before using the result in the next application. This
3263 -- function is strict in the starting value.
3264 foldrWithKey' :: (k -> a -> b -> b) -> b -> Map k a -> b
3265 foldrWithKey' f z = go z
3266 where
3267 go !z' Tip = z'
3268 go z' (Bin _ kx x l r) = go (f kx x (go z' r)) l
3269 {-# INLINE foldrWithKey' #-}
3270
3271 -- | /O(n)/. Fold the keys and values in the map using the given left-associative
3272 -- binary operator, such that
3273 -- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@.
3274 --
3275 -- For example,
3276 --
3277 -- > keys = reverse . foldlWithKey (\ks k x -> k:ks) []
3278 --
3279 -- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
3280 -- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)"
3281 foldlWithKey :: (a -> k -> b -> a) -> a -> Map k b -> a
3282 foldlWithKey f z = go z
3283 where
3284 go z' Tip = z'
3285 go z' (Bin _ kx x l r) = go (f (go z' l) kx x) r
3286 {-# INLINE foldlWithKey #-}
3287
3288 -- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is
3289 -- evaluated before using the result in the next application. This
3290 -- function is strict in the starting value.
3291 foldlWithKey' :: (a -> k -> b -> a) -> a -> Map k b -> a
3292 foldlWithKey' f z = go z
3293 where
3294 go !z' Tip = z'
3295 go z' (Bin _ kx x l r) = go (f (go z' l) kx x) r
3296 {-# INLINE foldlWithKey' #-}
3297
3298 -- | /O(n)/. Fold the keys and values in the map using the given monoid, such that
3299 --
3300 -- @'foldMapWithKey' f = 'Prelude.fold' . 'mapWithKey' f@
3301 --
3302 -- This can be an asymptotically faster than 'foldrWithKey' or 'foldlWithKey' for some monoids.
3303 --
3304 -- @since 0.5.4
3305 foldMapWithKey :: Monoid m => (k -> a -> m) -> Map k a -> m
3306 foldMapWithKey f = go
3307 where
3308 go Tip = mempty
3309 go (Bin 1 k v _ _) = f k v
3310 go (Bin _ k v l r) = go l `mappend` (f k v `mappend` go r)
3311 {-# INLINE foldMapWithKey #-}
3312
3313 {--------------------------------------------------------------------
3314 List variations
3315 --------------------------------------------------------------------}
3316 -- | /O(n)/.
3317 -- Return all elements of the map in the ascending order of their keys.
3318 -- Subject to list fusion.
3319 --
3320 -- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"]
3321 -- > elems empty == []
3322
3323 elems :: Map k a -> [a]
3324 elems = foldr (:) []
3325
3326 -- | /O(n)/. Return all keys of the map in ascending order. Subject to list
3327 -- fusion.
3328 --
3329 -- > keys (fromList [(5,"a"), (3,"b")]) == [3,5]
3330 -- > keys empty == []
3331
3332 keys :: Map k a -> [k]
3333 keys = foldrWithKey (\k _ ks -> k : ks) []
3334
3335 -- | /O(n)/. An alias for 'toAscList'. Return all key\/value pairs in the map
3336 -- in ascending key order. Subject to list fusion.
3337 --
3338 -- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
3339 -- > assocs empty == []
3340
3341 assocs :: Map k a -> [(k,a)]
3342 assocs m
3343 = toAscList m
3344
3345 -- | /O(n)/. The set of all keys of the map.
3346 --
3347 -- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.Set.fromList [3,5]
3348 -- > keysSet empty == Data.Set.empty
3349
3350 keysSet :: Map k a -> Set.Set k
3351 keysSet Tip = Set.Tip
3352 keysSet (Bin sz kx _ l r) = Set.Bin sz kx (keysSet l) (keysSet r)
3353
3354 -- | /O(n)/. Build a map from a set of keys and a function which for each key
3355 -- computes its value.
3356 --
3357 -- > fromSet (\k -> replicate k 'a') (Data.Set.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")]
3358 -- > fromSet undefined Data.Set.empty == empty
3359
3360 fromSet :: (k -> a) -> Set.Set k -> Map k a
3361 fromSet _ Set.Tip = Tip
3362 fromSet f (Set.Bin sz x l r) = Bin sz x (f x) (fromSet f l) (fromSet f r)
3363
3364 {--------------------------------------------------------------------
3365 Lists
3366 --------------------------------------------------------------------}
3367 #if __GLASGOW_HASKELL__ >= 708
3368 -- | @since 0.5.6.2
3369 instance (Ord k) => GHCExts.IsList (Map k v) where
3370 type Item (Map k v) = (k,v)
3371 fromList = fromList
3372 toList = toList
3373 #endif
3374
3375 -- | /O(n*log n)/. Build a map from a list of key\/value pairs. See also 'fromAscList'.
3376 -- If the list contains more than one value for the same key, the last value
3377 -- for the key is retained.
3378 --
3379 -- If the keys of the list are ordered, linear-time implementation is used,
3380 -- with the performance equal to 'fromDistinctAscList'.
3381 --
3382 -- > fromList [] == empty
3383 -- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
3384 -- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
3385
3386 -- For some reason, when 'singleton' is used in fromList or in
3387 -- create, it is not inlined, so we inline it manually.
3388 fromList :: Ord k => [(k,a)] -> Map k a
3389 fromList [] = Tip
3390 fromList [(kx, x)] = Bin 1 kx x Tip Tip
3391 fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = fromList' (Bin 1 kx0 x0 Tip Tip) xs0
3392 | otherwise = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0
3393 where
3394 not_ordered _ [] = False
3395 not_ordered kx ((ky,_) : _) = kx >= ky
3396 {-# INLINE not_ordered #-}
3397
3398 fromList' t0 xs = Foldable.foldl' ins t0 xs
3399 where ins t (k,x) = insert k x t
3400
3401 go !_ t [] = t
3402 go _ t [(kx, x)] = insertMax kx x t
3403 go s l xs@((kx, x) : xss) | not_ordered kx xss = fromList' l xs
3404 | otherwise = case create s xss of
3405 (r, ys, []) -> go (s `shiftL` 1) (link kx x l r) ys
3406 (r, _, ys) -> fromList' (link kx x l r) ys
3407
3408 -- The create is returning a triple (tree, xs, ys). Both xs and ys
3409 -- represent not yet processed elements and only one of them can be nonempty.
3410 -- If ys is nonempty, the keys in ys are not ordered with respect to tree
3411 -- and must be inserted using fromList'. Otherwise the keys have been
3412 -- ordered so far.
3413 create !_ [] = (Tip, [], [])
3414 create s xs@(xp : xss)
3415 | s == 1 = case xp of (kx, x) | not_ordered kx xss -> (Bin 1 kx x Tip Tip, [], xss)
3416 | otherwise -> (Bin 1 kx x Tip Tip, xss, [])
3417 | otherwise = case create (s `shiftR` 1) xs of
3418 res@(_, [], _) -> res
3419 (l, [(ky, y)], zs) -> (insertMax ky y l, [], zs)
3420 (l, ys@((ky, y):yss), _) | not_ordered ky yss -> (l, [], ys)
3421 | otherwise -> case create (s `shiftR` 1) yss of
3422 (r, zs, ws) -> (link ky y l r, zs, ws)
3423 #if __GLASGOW_HASKELL__
3424 {-# INLINABLE fromList #-}
3425 #endif
3426
3427 -- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
3428 --
3429 -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")]
3430 -- > fromListWith (++) [] == empty
3431
3432 fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a
3433 fromListWith f xs
3434 = fromListWithKey (\_ x y -> f x y) xs
3435 #if __GLASGOW_HASKELL__
3436 {-# INLINABLE fromListWith #-}
3437 #endif
3438
3439 -- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'.
3440 --
3441 -- > let f k a1 a2 = (show k) ++ a1 ++ a2
3442 -- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "3ab"), (5, "5a5ba")]
3443 -- > fromListWithKey f [] == empty
3444
3445 fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
3446 fromListWithKey f xs
3447 = Foldable.foldl' ins empty xs
3448 where
3449 ins t (k,x) = insertWithKey f k x t
3450 #if __GLASGOW_HASKELL__
3451 {-# INLINABLE fromListWithKey #-}
3452 #endif
3453
3454 -- | /O(n)/. Convert the map to a list of key\/value pairs. Subject to list fusion.
3455 --
3456 -- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
3457 -- > toList empty == []
3458
3459 toList :: Map k a -> [(k,a)]
3460 toList = toAscList
3461
3462 -- | /O(n)/. Convert the map to a list of key\/value pairs where the keys are
3463 -- in ascending order. Subject to list fusion.
3464 --
3465 -- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
3466
3467 toAscList :: Map k a -> [(k,a)]
3468 toAscList = foldrWithKey (\k x xs -> (k,x):xs) []
3469
3470 -- | /O(n)/. Convert the map to a list of key\/value pairs where the keys
3471 -- are in descending order. Subject to list fusion.
3472 --
3473 -- > toDescList (fromList [(5,"a"), (3,"b")]) == [(5,"a"), (3,"b")]
3474
3475 toDescList :: Map k a -> [(k,a)]
3476 toDescList = foldlWithKey (\xs k x -> (k,x):xs) []
3477
3478 -- List fusion for the list generating functions.
3479 #if __GLASGOW_HASKELL__
3480 -- The foldrFB and foldlFB are fold{r,l}WithKey equivalents, used for list fusion.
3481 -- They are important to convert unfused methods back, see mapFB in prelude.
3482 foldrFB :: (k -> a -> b -> b) -> b -> Map k a -> b
3483 foldrFB = foldrWithKey
3484 {-# INLINE[0] foldrFB #-}
3485 foldlFB :: (a -> k -> b -> a) -> a -> Map k b -> a
3486 foldlFB = foldlWithKey
3487 {-# INLINE[0] foldlFB #-}
3488
3489 -- Inline assocs and toList, so that we need to fuse only toAscList.
3490 {-# INLINE assocs #-}
3491 {-# INLINE toList #-}
3492
3493 -- The fusion is enabled up to phase 2 included. If it does not succeed,
3494 -- convert in phase 1 the expanded elems,keys,to{Asc,Desc}List calls back to
3495 -- elems,keys,to{Asc,Desc}List. In phase 0, we inline fold{lr}FB (which were
3496 -- used in a list fusion, otherwise it would go away in phase 1), and let compiler
3497 -- do whatever it wants with elems,keys,to{Asc,Desc}List -- it was forbidden to
3498 -- inline it before phase 0, otherwise the fusion rules would not fire at all.
3499 {-# NOINLINE[0] elems #-}
3500 {-# NOINLINE[0] keys #-}
3501 {-# NOINLINE[0] toAscList #-}
3502 {-# NOINLINE[0] toDescList #-}
3503 {-# RULES "Map.elems" [~1] forall m . elems m = build (\c n -> foldrFB (\_ x xs -> c x xs) n m) #-}
3504 {-# RULES "Map.elemsBack" [1] foldrFB (\_ x xs -> x : xs) [] = elems #-}
3505 {-# RULES "Map.keys" [~1] forall m . keys m = build (\c n -> foldrFB (\k _ xs -> c k xs) n m) #-}
3506 {-# RULES "Map.keysBack" [1] foldrFB (\k _ xs -> k : xs) [] = keys #-}
3507 {-# RULES "Map.toAscList" [~1] forall m . toAscList m = build (\c n -> foldrFB (\k x xs -> c (k,x) xs) n m) #-}
3508 {-# RULES "Map.toAscListBack" [1] foldrFB (\k x xs -> (k, x) : xs) [] = toAscList #-}
3509 {-# RULES "Map.toDescList" [~1] forall m . toDescList m = build (\c n -> foldlFB (\xs k x -> c (k,x) xs) n m) #-}
3510 {-# RULES "Map.toDescListBack" [1] foldlFB (\xs k x -> (k, x) : xs) [] = toDescList #-}
3511 #endif
3512
3513 {--------------------------------------------------------------------
3514 Building trees from ascending/descending lists can be done in linear time.
3515
3516 Note that if [xs] is ascending that:
3517 fromAscList xs == fromList xs
3518 fromAscListWith f xs == fromListWith f xs
3519 --------------------------------------------------------------------}
3520 -- | /O(n)/. Build a map from an ascending list in linear time.
3521 -- /The precondition (input list is ascending) is not checked./
3522 --
3523 -- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
3524 -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")]
3525 -- > valid (fromAscList [(3,"b"), (5,"a"), (5,"b")]) == True
3526 -- > valid (fromAscList [(5,"a"), (3,"b"), (5,"b")]) == False
3527
3528 fromAscList :: Eq k => [(k,a)] -> Map k a
3529 fromAscList xs
3530 = fromDistinctAscList (combineEq xs)
3531 where
3532 -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
3533 combineEq xs'
3534 = case xs' of
3535 [] -> []
3536 [x] -> [x]
3537 (x:xx) -> combineEq' x xx
3538
3539 combineEq' z [] = [z]
3540 combineEq' z@(kz,_) (x@(kx,xx):xs')
3541 | kx==kz = combineEq' (kx,xx) xs'
3542 | otherwise = z:combineEq' x xs'
3543 #if __GLASGOW_HASKELL__
3544 {-# INLINABLE fromAscList #-}
3545 #endif
3546
3547 -- | /O(n)/. Build a map from a descending list in linear time.
3548 -- /The precondition (input list is descending) is not checked./
3549 --
3550 -- > fromDescList [(5,"a"), (3,"b")] == fromList [(3, "b"), (5, "a")]
3551 -- > fromDescList [(5,"a"), (5,"b"), (3,"b")] == fromList [(3, "b"), (5, "b")]
3552 -- > valid (fromDescList [(5,"a"), (5,"b"), (3,"b")]) == True
3553 -- > valid (fromDescList [(5,"a"), (3,"b"), (5,"b")]) == False
3554 --
3555 -- @since 0.5.8
3556
3557 fromDescList :: Eq k => [(k,a)] -> Map k a
3558 fromDescList xs = fromDistinctDescList (combineEq xs)
3559 where
3560 -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
3561 combineEq xs'
3562 = case xs' of
3563 [] -> []
3564 [x] -> [x]
3565 (x:xx) -> combineEq' x xx
3566
3567 combineEq' z [] = [z]
3568 combineEq' z@(kz,_) (x@(kx,xx):xs')
3569 | kx==kz = combineEq' (kx,xx) xs'
3570 | otherwise = z:combineEq' x xs'
3571 #if __GLASGOW_HASKELL__
3572 {-# INLINABLE fromDescList #-}
3573 #endif
3574
3575 -- | /O(n)/. Build a map from an ascending list in linear time with a combining function for equal keys.
3576 -- /The precondition (input list is ascending) is not checked./
3577 --
3578 -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
3579 -- > valid (fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")]) == True
3580 -- > valid (fromAscListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False
3581
3582 fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a
3583 fromAscListWith f xs
3584 = fromAscListWithKey (\_ x y -> f x y) xs
3585 #if __GLASGOW_HASKELL__
3586 {-# INLINABLE fromAscListWith #-}
3587 #endif
3588
3589 -- | /O(n)/. Build a map from a descending list in linear time with a combining function for equal keys.
3590 -- /The precondition (input list is descending) is not checked./
3591 --
3592 -- > fromDescListWith (++) [(5,"a"), (5,"b"), (3,"b")] == fromList [(3, "b"), (5, "ba")]
3593 -- > valid (fromDescListWith (++) [(5,"a"), (5,"b"), (3,"b")]) == True
3594 -- > valid (fromDescListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False
3595 --
3596 -- @since 0.5.8
3597
3598 fromDescListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a
3599 fromDescListWith f xs
3600 = fromDescListWithKey (\_ x y -> f x y) xs
3601 #if __GLASGOW_HASKELL__
3602 {-# INLINABLE fromDescListWith #-}
3603 #endif
3604
3605 -- | /O(n)/. Build a map from an ascending list in linear time with a
3606 -- combining function for equal keys.
3607 -- /The precondition (input list is ascending) is not checked./
3608 --
3609 -- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2
3610 -- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")] == fromList [(3, "b"), (5, "5:b5:ba")]
3611 -- > valid (fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")]) == True
3612 -- > valid (fromAscListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False
3613
3614 fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
3615 fromAscListWithKey f xs
3616 = fromDistinctAscList (combineEq f xs)
3617 where
3618 -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
3619 combineEq _ xs'
3620 = case xs' of
3621 [] -> []
3622 [x] -> [x]
3623 (x:xx) -> combineEq' x xx
3624
3625 combineEq' z [] = [z]
3626 combineEq' z@(kz,zz) (x@(kx,xx):xs')
3627 | kx==kz = let yy = f kx xx zz in combineEq' (kx,yy) xs'
3628 | otherwise = z:combineEq' x xs'
3629 #if __GLASGOW_HASKELL__
3630 {-# INLINABLE fromAscListWithKey #-}
3631 #endif
3632
3633 -- | /O(n)/. Build a map from a descending list in linear time with a
3634 -- combining function for equal keys.
3635 -- /The precondition (input list is descending) is not checked./
3636 --
3637 -- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2
3638 -- > fromDescListWithKey f [(5,"a"), (5,"b"), (5,"b"), (3,"b")] == fromList [(3, "b"), (5, "5:b5:ba")]
3639 -- > valid (fromDescListWithKey f [(5,"a"), (5,"b"), (5,"b"), (3,"b")]) == True
3640 -- > valid (fromDescListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False
3641 fromDescListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
3642 fromDescListWithKey f xs
3643 = fromDistinctDescList (combineEq f xs)
3644 where
3645 -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
3646 combineEq _ xs'
3647 = case xs' of
3648 [] -> []
3649 [x] -> [x]
3650 (x:xx) -> combineEq' x xx
3651
3652 combineEq' z [] = [z]
3653 combineEq' z@(kz,zz) (x@(kx,xx):xs')
3654 | kx==kz = let yy = f kx xx zz in combineEq' (kx,yy) xs'
3655 | otherwise = z:combineEq' x xs'
3656 #if __GLASGOW_HASKELL__
3657 {-# INLINABLE fromDescListWithKey #-}
3658 #endif
3659
3660
3661 -- | /O(n)/. Build a map from an ascending list of distinct elements in linear time.
3662 -- /The precondition is not checked./
3663 --
3664 -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
3665 -- > valid (fromDistinctAscList [(3,"b"), (5,"a")]) == True
3666 -- > valid (fromDistinctAscList [(3,"b"), (5,"a"), (5,"b")]) == False
3667
3668 -- For some reason, when 'singleton' is used in fromDistinctAscList or in
3669 -- create, it is not inlined, so we inline it manually.
3670 fromDistinctAscList :: [(k,a)] -> Map k a
3671 fromDistinctAscList [] = Tip
3672 fromDistinctAscList ((kx0, x0) : xs0) = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0
3673 where
3674 go !_ t [] = t
3675 go s l ((kx, x) : xs) = case create s xs of
3676 (r :*: ys) -> let !t' = link kx x l r
3677 in go (s `shiftL` 1) t' ys
3678
3679 create !_ [] = (Tip :*: [])
3680 create s xs@(x' : xs')
3681 | s == 1 = case x' of (kx, x) -> (Bin 1 kx x Tip Tip :*: xs')
3682 | otherwise = case create (s `shiftR` 1) xs of
3683 res@(_ :*: []) -> res
3684 (l :*: (ky, y):ys) -> case create (s `shiftR` 1) ys of
3685 (r :*: zs) -> (link ky y l r :*: zs)
3686
3687 -- | /O(n)/. Build a map from a descending list of distinct elements in linear time.
3688 -- /The precondition is not checked./
3689 --
3690 -- > fromDistinctDescList [(5,"a"), (3,"b")] == fromList [(3, "b"), (5, "a")]
3691 -- > valid (fromDistinctDescList [(5,"a"), (3,"b")]) == True
3692 -- > valid (fromDistinctDescList [(5,"a"), (5,"b"), (3,"b")]) == False
3693 --
3694 -- @since 0.5.8
3695
3696 -- For some reason, when 'singleton' is used in fromDistinctDescList or in
3697 -- create, it is not inlined, so we inline it manually.
3698 fromDistinctDescList :: [(k,a)] -> Map k a
3699 fromDistinctDescList [] = Tip
3700 fromDistinctDescList ((kx0, x0) : xs0) = go (1 :: Int) (Bin 1 kx0 x0 Tip Tip) xs0
3701 where
3702 go !_ t [] = t
3703 go s r ((kx, x) : xs) = case create s xs of
3704 (l :*: ys) -> let !t' = link kx x l r
3705 in go (s `shiftL` 1) t' ys
3706
3707 create !_ [] = (Tip :*: [])
3708 create s xs@(x' : xs')
3709 | s == 1 = case x' of (kx, x) -> (Bin 1 kx x Tip Tip :*: xs')
3710 | otherwise = case create (s `shiftR` 1) xs of
3711 res@(_ :*: []) -> res
3712 (r :*: (ky, y):ys) -> case create (s `shiftR` 1) ys of
3713 (l :*: zs) -> (link ky y l r :*: zs)
3714
3715 {-
3716 -- Functions very similar to these were used to implement
3717 -- hedge union, intersection, and difference algorithms that we no
3718 -- longer use. These functions, however, seem likely to be useful
3719 -- in their own right, so I'm leaving them here in case we end up
3720 -- exporting them.
3721
3722 {--------------------------------------------------------------------
3723 [filterGt b t] filter all keys >[b] from tree [t]
3724 [filterLt b t] filter all keys <[b] from tree [t]
3725 --------------------------------------------------------------------}
3726 filterGt :: Ord k => k -> Map k v -> Map k v
3727 filterGt !_ Tip = Tip
3728 filterGt !b (Bin _ kx x l r) =
3729 case compare b kx of LT -> link kx x (filterGt b l) r
3730 EQ -> r
3731 GT -> filterGt b r
3732 #if __GLASGOW_HASKELL__
3733 {-# INLINABLE filterGt #-}
3734 #endif
3735
3736 filterLt :: Ord k => k -> Map k v -> Map k v
3737 filterLt !_ Tip = Tip
3738 filterLt !b (Bin _ kx x l r) =
3739 case compare kx b of LT -> link kx x l (filterLt b r)
3740 EQ -> l
3741 GT -> filterLt b l
3742 #if __GLASGOW_HASKELL__
3743 {-# INLINABLE filterLt #-}
3744 #endif
3745 -}
3746
3747 {--------------------------------------------------------------------
3748 Split
3749 --------------------------------------------------------------------}
3750 -- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@ where
3751 -- the keys in @map1@ are smaller than @k@ and the keys in @map2@ larger than @k@.
3752 -- Any key equal to @k@ is found in neither @map1@ nor @map2@.
3753 --
3754 -- > split 2 (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3,"b"), (5,"a")])
3755 -- > split 3 (fromList [(5,"a"), (3,"b")]) == (empty, singleton 5 "a")
3756 -- > split 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
3757 -- > split 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", empty)
3758 -- > split 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], empty)
3759
3760 split :: Ord k => k -> Map k a -> (Map k a,Map k a)
3761 split !k0 t0 = toPair $ go k0 t0
3762 where
3763 go k t =
3764 case t of
3765 Tip -> Tip :*: Tip
3766 Bin _ kx x l r -> case compare k kx of
3767 LT -> let (lt :*: gt) = go k l in lt :*: link kx x gt r
3768 GT -> let (lt :*: gt) = go k r in link kx x l lt :*: gt
3769 EQ -> (l :*: r)
3770 #if __GLASGOW_HASKELL__
3771 {-# INLINABLE split #-}
3772 #endif
3773
3774 -- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just
3775 -- like 'split' but also returns @'lookup' k map@.
3776 --
3777 -- > splitLookup 2 (fromList [(5,"a"), (3,"b")]) == (empty, Nothing, fromList [(3,"b"), (5,"a")])
3778 -- > splitLookup 3 (fromList [(5,"a"), (3,"b")]) == (empty, Just "b", singleton 5 "a")
3779 -- > splitLookup 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Nothing, singleton 5 "a")
3780 -- > splitLookup 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Just "a", empty)
3781 -- > splitLookup 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], Nothing, empty)
3782 splitLookup :: Ord k => k -> Map k a -> (Map k a,Maybe a,Map k a)
3783 splitLookup k0 m = case go k0 m of
3784 StrictTriple l mv r -> (l, mv, r)
3785 where
3786 go :: Ord k => k -> Map k a -> StrictTriple (Map k a) (Maybe a) (Map k a)
3787 go !k t =
3788 case t of
3789 Tip -> StrictTriple Tip Nothing Tip
3790 Bin _ kx x l r -> case compare k kx of
3791 LT -> let StrictTriple lt z gt = go k l
3792 !gt' = link kx x gt r
3793 in StrictTriple lt z gt'
3794 GT -> let StrictTriple lt z gt = go k r
3795 !lt' = link kx x l lt
3796 in StrictTriple lt' z gt
3797 EQ -> StrictTriple l (Just x) r
3798 #if __GLASGOW_HASKELL__
3799 {-# INLINABLE splitLookup #-}
3800 #endif
3801
3802 -- | A variant of 'splitLookup' that indicates only whether the
3803 -- key was present, rather than producing its value. This is used to
3804 -- implement 'intersection' to avoid allocating unnecessary 'Just'
3805 -- constructors.
3806 splitMember :: Ord k => k -> Map k a -> (Map k a,Bool,Map k a)
3807 splitMember k0 m = case go k0 m of
3808 StrictTriple l mv r -> (l, mv, r)
3809 where
3810 go :: Ord k => k -> Map k a -> StrictTriple (Map k a) Bool (Map k a)
3811 go !k t =
3812 case t of
3813 Tip -> StrictTriple Tip False Tip
3814 Bin _ kx x l r -> case compare k kx of
3815 LT -> let StrictTriple lt z gt = go k l
3816 !gt' = link kx x gt r
3817 in StrictTriple lt z gt'
3818 GT -> let StrictTriple lt z gt = go k r
3819 !lt' = link kx x l lt
3820 in StrictTriple lt' z gt
3821 EQ -> StrictTriple l True r
3822 #if __GLASGOW_HASKELL__
3823 {-# INLINABLE splitMember #-}
3824 #endif
3825
3826 data StrictTriple a b c = StrictTriple !a !b !c
3827
3828 {--------------------------------------------------------------------
3829 Utility functions that maintain the balance properties of the tree.
3830 All constructors assume that all values in [l] < [k] and all values
3831 in [r] > [k], and that [l] and [r] are valid trees.
3832
3833 In order of sophistication:
3834 [Bin sz k x l r] The type constructor.
3835 [bin k x l r] Maintains the correct size, assumes that both [l]
3836 and [r] are balanced with respect to each other.
3837 [balance k x l r] Restores the balance and size.
3838 Assumes that the original tree was balanced and
3839 that [l] or [r] has changed by at most one element.
3840 [link k x l r] Restores balance and size.
3841
3842 Furthermore, we can construct a new tree from two trees. Both operations
3843 assume that all values in [l] < all values in [r] and that [l] and [r]
3844 are valid:
3845 [glue l r] Glues [l] and [r] together. Assumes that [l] and
3846 [r] are already balanced with respect to each other.
3847 [link2 l r] Merges two trees and restores balance.
3848 --------------------------------------------------------------------}
3849
3850 {--------------------------------------------------------------------
3851 Link
3852 --------------------------------------------------------------------}
3853 link :: k -> a -> Map k a -> Map k a -> Map k a
3854 link kx x Tip r = insertMin kx x r
3855 link kx x l Tip = insertMax kx x l
3856 link kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz)
3857 | delta*sizeL < sizeR = balanceL kz z (link kx x l lz) rz
3858 | delta*sizeR < sizeL = balanceR ky y ly (link kx x ry r)
3859 | otherwise = bin kx x l r
3860
3861
3862 -- insertMin and insertMax don't perform potentially expensive comparisons.
3863 insertMax,insertMin :: k -> a -> Map k a -> Map k a
3864 insertMax kx x t
3865 = case t of
3866 Tip -> singleton kx x
3867 Bin _ ky y l r