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