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