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