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