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