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