92bfc1dcc32ef40f068d32275c327c5d626688b9
[packages/containers.git] / Data / Set / 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 #endif
13
14 #include "containers.h"
15
16 -----------------------------------------------------------------------------
17 -- |
18 -- Module : Data.Set.Base
19 -- Copyright : (c) Daan Leijen 2002
20 -- License : BSD-style
21 -- Maintainer : libraries@haskell.org
22 -- Stability : provisional
23 -- Portability : portable
24 --
25 -- An efficient implementation of sets.
26 --
27 -- These modules are intended to be imported qualified, to avoid name
28 -- clashes with Prelude functions, e.g.
29 --
30 -- > import Data.Set (Set)
31 -- > import qualified Data.Set as Set
32 --
33 -- The implementation of 'Set' is based on /size balanced/ binary trees (or
34 -- trees of /bounded balance/) as described by:
35 --
36 -- * Stephen Adams, \"/Efficient sets: a balancing act/\",
37 -- Journal of Functional Programming 3(4):553-562, October 1993,
38 -- <http://www.swiss.ai.mit.edu/~adams/BB/>.
39 --
40 -- * J. Nievergelt and E.M. Reingold,
41 -- \"/Binary search trees of bounded balance/\",
42 -- SIAM journal of computing 2(1), March 1973.
43 --
44 -- Note that the implementation is /left-biased/ -- the elements of a
45 -- first argument are always preferred to the second, for example in
46 -- 'union' or 'insert'. Of course, left-biasing can only be observed
47 -- when equality is an equivalence relation instead of structural
48 -- equality.
49 --
50 -- /Warning/: The size of the set must not exceed @maxBound::Int@. Violation of
51 -- this condition is not detected and if the size limit is exceeded, its
52 -- behaviour is undefined.
53 -----------------------------------------------------------------------------
54
55 -- [Note: Using INLINABLE]
56 -- ~~~~~~~~~~~~~~~~~~~~~~~
57 -- It is crucial to the performance that the functions specialize on the Ord
58 -- type when possible. GHC 7.0 and higher does this by itself when it sees th
59 -- unfolding of a function -- that is why all public functions are marked
60 -- INLINABLE (that exposes the unfolding).
61
62
63 -- [Note: Using INLINE]
64 -- ~~~~~~~~~~~~~~~~~~~~
65 -- For other compilers and GHC pre 7.0, we mark some of the functions INLINE.
66 -- We mark the functions that just navigate down the tree (lookup, insert,
67 -- delete and similar). That navigation code gets inlined and thus specialized
68 -- when possible. There is a price to pay -- code growth. The code INLINED is
69 -- therefore only the tree navigation, all the real work (rebalancing) is not
70 -- INLINED by using a NOINLINE.
71 --
72 -- All methods marked INLINE have to be nonrecursive -- a 'go' function doing
73 -- the real work is provided.
74
75
76 -- [Note: Type of local 'go' function]
77 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
78 -- If the local 'go' function uses an Ord class, it sometimes heap-allocates
79 -- the Ord dictionary when the 'go' function does not have explicit type.
80 -- In that case we give 'go' explicit type. But this slightly decrease
81 -- performance, as the resulting 'go' function can float out to top level.
82
83
84 -- [Note: Local 'go' functions and capturing]
85 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
86 -- As opposed to IntSet, when 'go' function captures an argument, increased
87 -- heap-allocation can occur: sometimes in a polymorphic function, the 'go'
88 -- floats out of its enclosing function and then it heap-allocates the
89 -- dictionary and the argument. Maybe it floats out too late and strictness
90 -- analyzer cannot see that these could be passed on stack.
91 --
92 -- For example, change 'member' so that its local 'go' function is not passing
93 -- argument x and then look at the resulting code for hedgeInt.
94
95
96 -- [Note: Order of constructors]
97 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
98 -- The order of constructors of Set matters when considering performance.
99 -- Currently in GHC 7.0, when type has 2 constructors, a forward conditional
100 -- jump is made when successfully matching second constructor. Successful match
101 -- of first constructor results in the forward jump not taken.
102 -- On GHC 7.0, reordering constructors from Tip | Bin to Bin | Tip
103 -- improves the benchmark by up to 10% on x86.
104
105 module Data.Set.Base (
106 -- * Set type
107 Set(..) -- instance Eq,Ord,Show,Read,Data,Typeable
108
109 -- * Operators
110 , (\\)
111
112 -- * Query
113 , null
114 , size
115 , member
116 , notMember
117 , lookupLT
118 , lookupGT
119 , lookupLE
120 , lookupGE
121 , isSubsetOf
122 , isProperSubsetOf
123
124 -- * Construction
125 , empty
126 , singleton
127 , insert
128 , delete
129
130 -- * Combine
131 , union
132 , unions
133 , difference
134 , intersection
135
136 -- * Filter
137 , filter
138 , partition
139 , split
140 , splitMember
141 , splitRoot
142
143 -- * Indexed
144 , lookupIndex
145 , findIndex
146 , elemAt
147 , deleteAt
148
149 -- * Map
150 , map
151 , mapMonotonic
152
153 -- * Folds
154 , foldr
155 , foldl
156 -- ** Strict folds
157 , foldr'
158 , foldl'
159 -- ** Legacy folds
160 , fold
161
162 -- * Min\/Max
163 , findMin
164 , findMax
165 , deleteMin
166 , deleteMax
167 , deleteFindMin
168 , deleteFindMax
169 , maxView
170 , minView
171
172 -- * Conversion
173
174 -- ** List
175 , elems
176 , toList
177 , fromList
178
179 -- ** Ordered list
180 , toAscList
181 , toDescList
182 , fromAscList
183 , fromDistinctAscList
184
185 -- * Debugging
186 , showTree
187 , showTreeWith
188 , valid
189
190 -- Internals (for testing)
191 , bin
192 , balanced
193 , link
194 , merge
195 ) where
196
197 import Prelude hiding (filter,foldl,foldr,null,map)
198 import qualified Data.List as List
199 import Data.Bits (shiftL, shiftR)
200 #if !MIN_VERSION_base(4,8,0)
201 import Data.Monoid (Monoid(..))
202 #endif
203 #if MIN_VERSION_base(4,9,0)
204 import Data.Semigroup (Semigroup((<>), stimes), stimesIdempotentMonoid)
205 #endif
206 import qualified Data.Foldable as Foldable
207 import Data.Typeable
208 import Control.DeepSeq (NFData(rnf))
209
210 import Data.Utils.StrictFold
211 import Data.Utils.StrictPair
212
213 #if __GLASGOW_HASKELL__
214 import GHC.Exts ( build )
215 #if __GLASGOW_HASKELL__ >= 708
216 import qualified GHC.Exts as GHCExts
217 #endif
218 import Text.Read
219 import Data.Data
220 #endif
221
222
223 {--------------------------------------------------------------------
224 Operators
225 --------------------------------------------------------------------}
226 infixl 9 \\ --
227
228 -- | /O(n+m)/. See 'difference'.
229 (\\) :: Ord a => Set a -> Set a -> Set a
230 m1 \\ m2 = difference m1 m2
231 #if __GLASGOW_HASKELL__
232 {-# INLINABLE (\\) #-}
233 #endif
234
235 {--------------------------------------------------------------------
236 Sets are size balanced trees
237 --------------------------------------------------------------------}
238 -- | A set of values @a@.
239
240 -- See Note: Order of constructors
241 data Set a = Bin {-# UNPACK #-} !Size !a !(Set a) !(Set a)
242 | Tip
243
244 type Size = Int
245
246 #if __GLASGOW_HASKELL__ >= 708
247 type role Set nominal
248 #endif
249
250 instance Ord a => Monoid (Set a) where
251 mempty = empty
252 mconcat = unions
253 #if !(MIN_VERSION_base(4,9,0))
254 mappend = union
255 #else
256 mappend = (<>)
257
258 instance Ord a => Semigroup (Set a) where
259 (<>) = union
260 stimes = stimesIdempotentMonoid
261 #endif
262
263
264 instance Foldable.Foldable Set where
265 fold = go
266 where go Tip = mempty
267 go (Bin 1 k _ _) = k
268 go (Bin _ k l r) = go l `mappend` (k `mappend` go r)
269 {-# INLINABLE fold #-}
270 foldr = foldr
271 {-# INLINE foldr #-}
272 foldl = foldl
273 {-# INLINE foldl #-}
274 foldMap f t = go t
275 where go Tip = mempty
276 go (Bin 1 k _ _) = f k
277 go (Bin _ k l r) = go l `mappend` (f k `mappend` go r)
278 {-# INLINE foldMap #-}
279
280 #if MIN_VERSION_base(4,6,0)
281 foldl' = foldl'
282 {-# INLINE foldl' #-}
283 foldr' = foldr'
284 {-# INLINE foldr' #-}
285 #endif
286 #if MIN_VERSION_base(4,8,0)
287 length = size
288 {-# INLINE length #-}
289 null = null
290 {-# INLINE null #-}
291 toList = toList
292 {-# INLINE toList #-}
293 elem = go
294 where go !_ Tip = False
295 go x (Bin _ y l r) = x == y || go x l || go x r
296 {-# INLINABLE elem #-}
297 minimum = findMin
298 {-# INLINE minimum #-}
299 maximum = findMax
300 {-# INLINE maximum #-}
301 sum = foldl' (+) 0
302 {-# INLINABLE sum #-}
303 product = foldl' (*) 1
304 {-# INLINABLE product #-}
305 #endif
306
307
308 #if __GLASGOW_HASKELL__
309
310 {--------------------------------------------------------------------
311 A Data instance
312 --------------------------------------------------------------------}
313
314 -- This instance preserves data abstraction at the cost of inefficiency.
315 -- We provide limited reflection services for the sake of data abstraction.
316
317 instance (Data a, Ord a) => Data (Set a) where
318 gfoldl f z set = z fromList `f` (toList set)
319 toConstr _ = fromListConstr
320 gunfold k z c = case constrIndex c of
321 1 -> k (z fromList)
322 _ -> error "gunfold"
323 dataTypeOf _ = setDataType
324 dataCast1 f = gcast1 f
325
326 fromListConstr :: Constr
327 fromListConstr = mkConstr setDataType "fromList" [] Prefix
328
329 setDataType :: DataType
330 setDataType = mkDataType "Data.Set.Base.Set" [fromListConstr]
331
332 #endif
333
334 {--------------------------------------------------------------------
335 Query
336 --------------------------------------------------------------------}
337 -- | /O(1)/. Is this the empty set?
338 null :: Set a -> Bool
339 null Tip = True
340 null (Bin {}) = False
341 {-# INLINE null #-}
342
343 -- | /O(1)/. The number of elements in the set.
344 size :: Set a -> Int
345 size Tip = 0
346 size (Bin sz _ _ _) = sz
347 {-# INLINE size #-}
348
349 -- | /O(log n)/. Is the element in the set?
350 member :: Ord a => a -> Set a -> Bool
351 member = go
352 where
353 go !_ Tip = False
354 go x (Bin _ y l r) = case compare x y of
355 LT -> go x l
356 GT -> go x r
357 EQ -> True
358 #if __GLASGOW_HASKELL__
359 {-# INLINABLE member #-}
360 #else
361 {-# INLINE member #-}
362 #endif
363
364 -- | /O(log n)/. Is the element not in the set?
365 notMember :: Ord a => a -> Set a -> Bool
366 notMember a t = not $ member a t
367 #if __GLASGOW_HASKELL__
368 {-# INLINABLE notMember #-}
369 #else
370 {-# INLINE notMember #-}
371 #endif
372
373 -- | /O(log n)/. Find largest element smaller than the given one.
374 --
375 -- > lookupLT 3 (fromList [3, 5]) == Nothing
376 -- > lookupLT 5 (fromList [3, 5]) == Just 3
377 lookupLT :: Ord a => a -> Set a -> Maybe a
378 lookupLT = goNothing
379 where
380 goNothing !_ Tip = Nothing
381 goNothing x (Bin _ y l r) | x <= y = goNothing x l
382 | otherwise = goJust x y r
383
384 goJust !_ best Tip = Just best
385 goJust x best (Bin _ y l r) | x <= y = goJust x best l
386 | otherwise = goJust x y r
387 #if __GLASGOW_HASKELL__
388 {-# INLINABLE lookupLT #-}
389 #else
390 {-# INLINE lookupLT #-}
391 #endif
392
393 -- | /O(log n)/. Find smallest element greater than the given one.
394 --
395 -- > lookupGT 4 (fromList [3, 5]) == Just 5
396 -- > lookupGT 5 (fromList [3, 5]) == Nothing
397 lookupGT :: Ord a => a -> Set a -> Maybe a
398 lookupGT = goNothing
399 where
400 goNothing !_ Tip = Nothing
401 goNothing x (Bin _ y l r) | x < y = goJust x y l
402 | otherwise = goNothing x r
403
404 goJust !_ best Tip = Just best
405 goJust x best (Bin _ y l r) | x < y = goJust x y l
406 | otherwise = goJust x best r
407 #if __GLASGOW_HASKELL__
408 {-# INLINABLE lookupGT #-}
409 #else
410 {-# INLINE lookupGT #-}
411 #endif
412
413 -- | /O(log n)/. Find largest element smaller or equal to the given one.
414 --
415 -- > lookupLE 2 (fromList [3, 5]) == Nothing
416 -- > lookupLE 4 (fromList [3, 5]) == Just 3
417 -- > lookupLE 5 (fromList [3, 5]) == Just 5
418 lookupLE :: Ord a => a -> Set a -> Maybe a
419 lookupLE = goNothing
420 where
421 goNothing !_ Tip = Nothing
422 goNothing x (Bin _ y l r) = case compare x y of LT -> goNothing x l
423 EQ -> Just y
424 GT -> goJust x y r
425
426 goJust !_ best Tip = Just best
427 goJust x best (Bin _ y l r) = case compare x y of LT -> goJust x best l
428 EQ -> Just y
429 GT -> goJust x y r
430 #if __GLASGOW_HASKELL__
431 {-# INLINABLE lookupLE #-}
432 #else
433 {-# INLINE lookupLE #-}
434 #endif
435
436 -- | /O(log n)/. Find smallest element greater or equal to the given one.
437 --
438 -- > lookupGE 3 (fromList [3, 5]) == Just 3
439 -- > lookupGE 4 (fromList [3, 5]) == Just 5
440 -- > lookupGE 6 (fromList [3, 5]) == Nothing
441 lookupGE :: Ord a => a -> Set a -> Maybe a
442 lookupGE = goNothing
443 where
444 goNothing !_ Tip = Nothing
445 goNothing x (Bin _ y l r) = case compare x y of LT -> goJust x y l
446 EQ -> Just y
447 GT -> goNothing x r
448
449 goJust !_ best Tip = Just best
450 goJust x best (Bin _ y l r) = case compare x y of LT -> goJust x y l
451 EQ -> Just y
452 GT -> goJust x best r
453 #if __GLASGOW_HASKELL__
454 {-# INLINABLE lookupGE #-}
455 #else
456 {-# INLINE lookupGE #-}
457 #endif
458
459 {--------------------------------------------------------------------
460 Construction
461 --------------------------------------------------------------------}
462 -- | /O(1)/. The empty set.
463 empty :: Set a
464 empty = Tip
465 {-# INLINE empty #-}
466
467 -- | /O(1)/. Create a singleton set.
468 singleton :: a -> Set a
469 singleton x = Bin 1 x Tip Tip
470 {-# INLINE singleton #-}
471
472 {--------------------------------------------------------------------
473 Insertion, Deletion
474 --------------------------------------------------------------------}
475 -- | /O(log n)/. Insert an element in a set.
476 -- If the set already contains an element equal to the given value,
477 -- it is replaced with the new value.
478
479 -- See Note: Type of local 'go' function
480 insert :: Ord a => a -> Set a -> Set a
481 insert = go
482 where
483 go :: Ord a => a -> Set a -> Set a
484 go !x Tip = singleton x
485 go x (Bin sz y l r) = case compare x y of
486 LT -> balanceL y (go x l) r
487 GT -> balanceR y l (go x r)
488 EQ -> Bin sz x l r
489 #if __GLASGOW_HASKELL__
490 {-# INLINABLE insert #-}
491 #else
492 {-# INLINE insert #-}
493 #endif
494
495 -- Insert an element to the set only if it is not in the set.
496 -- Used by `union`.
497
498 -- See Note: Type of local 'go' function
499 insertR :: Ord a => a -> Set a -> Set a
500 insertR = go
501 where
502 go :: Ord a => a -> Set a -> Set a
503 go !x Tip = singleton x
504 go x t@(Bin _ y l r) = case compare x y of
505 LT -> balanceL y (go x l) r
506 GT -> balanceR y l (go x r)
507 EQ -> t
508 #if __GLASGOW_HASKELL__
509 {-# INLINABLE insertR #-}
510 #else
511 {-# INLINE insertR #-}
512 #endif
513
514 -- | /O(log n)/. Delete an element from a set.
515
516 -- See Note: Type of local 'go' function
517 delete :: Ord a => a -> Set a -> Set a
518 delete = go
519 where
520 go :: Ord a => a -> Set a -> Set a
521 go !_ Tip = Tip
522 go x (Bin _ y l r) = case compare x y of
523 LT -> balanceR y (go x l) r
524 GT -> balanceL y l (go x r)
525 EQ -> glue l r
526 #if __GLASGOW_HASKELL__
527 {-# INLINABLE delete #-}
528 #else
529 {-# INLINE delete #-}
530 #endif
531
532 {--------------------------------------------------------------------
533 Subset
534 --------------------------------------------------------------------}
535 -- | /O(n+m)/. Is this a proper subset? (ie. a subset but not equal).
536 isProperSubsetOf :: Ord a => Set a -> Set a -> Bool
537 isProperSubsetOf s1 s2
538 = (size s1 < size s2) && (isSubsetOf s1 s2)
539 #if __GLASGOW_HASKELL__
540 {-# INLINABLE isProperSubsetOf #-}
541 #endif
542
543
544 -- | /O(n+m)/. Is this a subset?
545 -- @(s1 `isSubsetOf` s2)@ tells whether @s1@ is a subset of @s2@.
546 isSubsetOf :: Ord a => Set a -> Set a -> Bool
547 isSubsetOf t1 t2
548 = (size t1 <= size t2) && (isSubsetOfX t1 t2)
549 #if __GLASGOW_HASKELL__
550 {-# INLINABLE isSubsetOf #-}
551 #endif
552
553 isSubsetOfX :: Ord a => Set a -> Set a -> Bool
554 isSubsetOfX Tip _ = True
555 isSubsetOfX _ Tip = False
556 isSubsetOfX (Bin _ x l r) t
557 = found && isSubsetOfX l lt && isSubsetOfX r gt
558 where
559 (lt,found,gt) = splitMember x t
560 #if __GLASGOW_HASKELL__
561 {-# INLINABLE isSubsetOfX #-}
562 #endif
563
564
565 {--------------------------------------------------------------------
566 Minimal, Maximal
567 --------------------------------------------------------------------}
568 -- | /O(log n)/. The minimal element of a set.
569 findMin :: Set a -> a
570 findMin (Bin _ x Tip _) = x
571 findMin (Bin _ _ l _) = findMin l
572 findMin Tip = error "Set.findMin: empty set has no minimal element"
573
574 -- | /O(log n)/. The maximal element of a set.
575 findMax :: Set a -> a
576 findMax (Bin _ x _ Tip) = x
577 findMax (Bin _ _ _ r) = findMax r
578 findMax Tip = error "Set.findMax: empty set has no maximal element"
579
580 -- | /O(log n)/. Delete the minimal element. Returns an empty set if the set is empty.
581 deleteMin :: Set a -> Set a
582 deleteMin (Bin _ _ Tip r) = r
583 deleteMin (Bin _ x l r) = balanceR x (deleteMin l) r
584 deleteMin Tip = Tip
585
586 -- | /O(log n)/. Delete the maximal element. Returns an empty set if the set is empty.
587 deleteMax :: Set a -> Set a
588 deleteMax (Bin _ _ l Tip) = l
589 deleteMax (Bin _ x l r) = balanceL x l (deleteMax r)
590 deleteMax Tip = Tip
591
592 {--------------------------------------------------------------------
593 Union.
594 --------------------------------------------------------------------}
595 -- | The union of a list of sets: (@'unions' == 'foldl' 'union' 'empty'@).
596 unions :: Ord a => [Set a] -> Set a
597 unions = foldlStrict union empty
598 #if __GLASGOW_HASKELL__
599 {-# INLINABLE unions #-}
600 #endif
601
602 -- | /O(n+m)/. The union of two sets, preferring the first set when
603 -- equal elements are encountered.
604 -- The implementation uses the efficient /hedge-union/ algorithm.
605 union :: Ord a => Set a -> Set a -> Set a
606 union Tip t2 = t2
607 union t1 Tip = t1
608 union t1 t2 = hedgeUnion NothingS NothingS t1 t2
609 #if __GLASGOW_HASKELL__
610 {-# INLINABLE union #-}
611 #endif
612
613 hedgeUnion :: Ord a => MaybeS a -> MaybeS a -> Set a -> Set a -> Set a
614 hedgeUnion _ _ t1 Tip = t1
615 hedgeUnion blo bhi Tip (Bin _ x l r) = link x (filterGt blo l) (filterLt bhi r)
616 hedgeUnion _ _ t1 (Bin _ x Tip Tip) = insertR x t1 -- According to benchmarks, this special case increases
617 -- performance up to 30%. It does not help in difference or intersection.
618 hedgeUnion blo bhi (Bin _ x l r) t2 = link x (hedgeUnion blo bmi l (trim blo bmi t2))
619 (hedgeUnion bmi bhi r (trim bmi bhi t2))
620 where bmi = JustS x
621 #if __GLASGOW_HASKELL__
622 {-# INLINABLE hedgeUnion #-}
623 #endif
624
625 {--------------------------------------------------------------------
626 Difference
627 --------------------------------------------------------------------}
628 -- | /O(n+m)/. Difference of two sets.
629 -- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
630 difference :: Ord a => Set a -> Set a -> Set a
631 difference Tip _ = Tip
632 difference t1 Tip = t1
633 difference t1 t2 = hedgeDiff NothingS NothingS t1 t2
634 #if __GLASGOW_HASKELL__
635 {-# INLINABLE difference #-}
636 #endif
637
638 hedgeDiff :: Ord a => MaybeS a -> MaybeS a -> Set a -> Set a -> Set a
639 hedgeDiff _ _ Tip _ = Tip
640 hedgeDiff blo bhi (Bin _ x l r) Tip = link x (filterGt blo l) (filterLt bhi r)
641 hedgeDiff blo bhi t (Bin _ x l r) = merge (hedgeDiff blo bmi (trim blo bmi t) l)
642 (hedgeDiff bmi bhi (trim bmi bhi t) r)
643 where bmi = JustS x
644 #if __GLASGOW_HASKELL__
645 {-# INLINABLE hedgeDiff #-}
646 #endif
647
648 {--------------------------------------------------------------------
649 Intersection
650 --------------------------------------------------------------------}
651 -- | /O(n+m)/. The intersection of two sets. The implementation uses an
652 -- efficient /hedge/ algorithm comparable with /hedge-union/. Elements of the
653 -- result come from the first set, so for example
654 --
655 -- > import qualified Data.Set as S
656 -- > data AB = A | B deriving Show
657 -- > instance Ord AB where compare _ _ = EQ
658 -- > instance Eq AB where _ == _ = True
659 -- > main = print (S.singleton A `S.intersection` S.singleton B,
660 -- > S.singleton B `S.intersection` S.singleton A)
661 --
662 -- prints @(fromList [A],fromList [B])@.
663 intersection :: Ord a => Set a -> Set a -> Set a
664 intersection Tip _ = Tip
665 intersection _ Tip = Tip
666 intersection t1 t2 = hedgeInt NothingS NothingS t1 t2
667 #if __GLASGOW_HASKELL__
668 {-# INLINABLE intersection #-}
669 #endif
670
671 hedgeInt :: Ord a => MaybeS a -> MaybeS a -> Set a -> Set a -> Set a
672 hedgeInt _ _ _ Tip = Tip
673 hedgeInt _ _ Tip _ = Tip
674 hedgeInt blo bhi (Bin _ x l r) t2 = let l' = hedgeInt blo bmi l (trim blo bmi t2)
675 r' = hedgeInt bmi bhi r (trim bmi bhi t2)
676 in if x `member` t2 then link x l' r' else merge l' r'
677 where bmi = JustS x
678 #if __GLASGOW_HASKELL__
679 {-# INLINABLE hedgeInt #-}
680 #endif
681
682 {--------------------------------------------------------------------
683 Filter and partition
684 --------------------------------------------------------------------}
685 -- | /O(n)/. Filter all elements that satisfy the predicate.
686 filter :: (a -> Bool) -> Set a -> Set a
687 filter _ Tip = Tip
688 filter p (Bin _ x l r)
689 | p x = link x (filter p l) (filter p r)
690 | otherwise = merge (filter p l) (filter p r)
691
692 -- | /O(n)/. Partition the set into two sets, one with all elements that satisfy
693 -- the predicate and one with all elements that don't satisfy the predicate.
694 -- See also 'split'.
695 partition :: (a -> Bool) -> Set a -> (Set a,Set a)
696 partition p0 t0 = toPair $ go p0 t0
697 where
698 go _ Tip = (Tip :*: Tip)
699 go p (Bin _ x l r) = case (go p l, go p r) of
700 ((l1 :*: l2), (r1 :*: r2))
701 | p x -> link x l1 r1 :*: merge l2 r2
702 | otherwise -> merge l1 r1 :*: link x l2 r2
703
704 {----------------------------------------------------------------------
705 Map
706 ----------------------------------------------------------------------}
707
708 -- | /O(n*log n)/.
709 -- @'map' f s@ is the set obtained by applying @f@ to each element of @s@.
710 --
711 -- It's worth noting that the size of the result may be smaller if,
712 -- for some @(x,y)@, @x \/= y && f x == f y@
713
714 map :: Ord b => (a->b) -> Set a -> Set b
715 map f = fromList . List.map f . toList
716 #if __GLASGOW_HASKELL__
717 {-# INLINABLE map #-}
718 #endif
719
720 -- | /O(n)/. The
721 --
722 -- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is monotonic.
723 -- /The precondition is not checked./
724 -- Semi-formally, we have:
725 --
726 -- > and [x < y ==> f x < f y | x <- ls, y <- ls]
727 -- > ==> mapMonotonic f s == map f s
728 -- > where ls = toList s
729
730 mapMonotonic :: (a->b) -> Set a -> Set b
731 mapMonotonic _ Tip = Tip
732 mapMonotonic f (Bin sz x l r) = Bin sz (f x) (mapMonotonic f l) (mapMonotonic f r)
733
734 {--------------------------------------------------------------------
735 Fold
736 --------------------------------------------------------------------}
737 -- | /O(n)/. Fold the elements in the set using the given right-associative
738 -- binary operator. This function is an equivalent of 'foldr' and is present
739 -- for compatibility only.
740 --
741 -- /Please note that fold will be deprecated in the future and removed./
742 fold :: (a -> b -> b) -> b -> Set a -> b
743 fold = foldr
744 {-# INLINE fold #-}
745
746 -- | /O(n)/. Fold the elements in the set using the given right-associative
747 -- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'toAscList'@.
748 --
749 -- For example,
750 --
751 -- > toAscList set = foldr (:) [] set
752 foldr :: (a -> b -> b) -> b -> Set a -> b
753 foldr f z = go z
754 where
755 go z' Tip = z'
756 go z' (Bin _ x l r) = go (f x (go z' r)) l
757 {-# INLINE foldr #-}
758
759 -- | /O(n)/. A strict version of 'foldr'. Each application of the operator is
760 -- evaluated before using the result in the next application. This
761 -- function is strict in the starting value.
762 foldr' :: (a -> b -> b) -> b -> Set a -> b
763 foldr' f z = go z
764 where
765 go !z' Tip = z'
766 go z' (Bin _ x l r) = go (f x (go z' r)) l
767 {-# INLINE foldr' #-}
768
769 -- | /O(n)/. Fold the elements in the set using the given left-associative
770 -- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'toAscList'@.
771 --
772 -- For example,
773 --
774 -- > toDescList set = foldl (flip (:)) [] set
775 foldl :: (a -> b -> a) -> a -> Set b -> a
776 foldl f z = go z
777 where
778 go z' Tip = z'
779 go z' (Bin _ x l r) = go (f (go z' l) x) r
780 {-# INLINE foldl #-}
781
782 -- | /O(n)/. A strict version of 'foldl'. Each application of the operator is
783 -- evaluated before using the result in the next application. This
784 -- function is strict in the starting value.
785 foldl' :: (a -> b -> a) -> a -> Set b -> a
786 foldl' f z = go z
787 where
788 go !z' Tip = z'
789 go z' (Bin _ x l r) = go (f (go z' l) x) r
790 {-# INLINE foldl' #-}
791
792 {--------------------------------------------------------------------
793 List variations
794 --------------------------------------------------------------------}
795 -- | /O(n)/. An alias of 'toAscList'. The elements of a set in ascending order.
796 -- Subject to list fusion.
797 elems :: Set a -> [a]
798 elems = toAscList
799
800 {--------------------------------------------------------------------
801 Lists
802 --------------------------------------------------------------------}
803 #if __GLASGOW_HASKELL__ >= 708
804 instance (Ord a) => GHCExts.IsList (Set a) where
805 type Item (Set a) = a
806 fromList = fromList
807 toList = toList
808 #endif
809
810 -- | /O(n)/. Convert the set to a list of elements. Subject to list fusion.
811 toList :: Set a -> [a]
812 toList = toAscList
813
814 -- | /O(n)/. Convert the set to an ascending list of elements. Subject to list fusion.
815 toAscList :: Set a -> [a]
816 toAscList = foldr (:) []
817
818 -- | /O(n)/. Convert the set to a descending list of elements. Subject to list
819 -- fusion.
820 toDescList :: Set a -> [a]
821 toDescList = foldl (flip (:)) []
822
823 -- List fusion for the list generating functions.
824 #if __GLASGOW_HASKELL__
825 -- The foldrFB and foldlFB are foldr and foldl equivalents, used for list fusion.
826 -- They are important to convert unfused to{Asc,Desc}List back, see mapFB in prelude.
827 foldrFB :: (a -> b -> b) -> b -> Set a -> b
828 foldrFB = foldr
829 {-# INLINE[0] foldrFB #-}
830 foldlFB :: (a -> b -> a) -> a -> Set b -> a
831 foldlFB = foldl
832 {-# INLINE[0] foldlFB #-}
833
834 -- Inline elems and toList, so that we need to fuse only toAscList.
835 {-# INLINE elems #-}
836 {-# INLINE toList #-}
837
838 -- The fusion is enabled up to phase 2 included. If it does not succeed,
839 -- convert in phase 1 the expanded to{Asc,Desc}List calls back to
840 -- to{Asc,Desc}List. In phase 0, we inline fold{lr}FB (which were used in
841 -- a list fusion, otherwise it would go away in phase 1), and let compiler do
842 -- whatever it wants with to{Asc,Desc}List -- it was forbidden to inline it
843 -- before phase 0, otherwise the fusion rules would not fire at all.
844 {-# NOINLINE[0] toAscList #-}
845 {-# NOINLINE[0] toDescList #-}
846 {-# RULES "Set.toAscList" [~1] forall s . toAscList s = build (\c n -> foldrFB c n s) #-}
847 {-# RULES "Set.toAscListBack" [1] foldrFB (:) [] = toAscList #-}
848 {-# RULES "Set.toDescList" [~1] forall s . toDescList s = build (\c n -> foldlFB (\xs x -> c x xs) n s) #-}
849 {-# RULES "Set.toDescListBack" [1] foldlFB (\xs x -> x : xs) [] = toDescList #-}
850 #endif
851
852 -- | /O(n*log n)/. Create a set from a list of elements.
853 --
854 -- If the elements are ordered, a linear-time implementation is used,
855 -- with the performance equal to 'fromDistinctAscList'.
856
857 -- For some reason, when 'singleton' is used in fromList or in
858 -- create, it is not inlined, so we inline it manually.
859 fromList :: Ord a => [a] -> Set a
860 fromList [] = Tip
861 fromList [x] = Bin 1 x Tip Tip
862 fromList (x0 : xs0) | not_ordered x0 xs0 = fromList' (Bin 1 x0 Tip Tip) xs0
863 | otherwise = go (1::Int) (Bin 1 x0 Tip Tip) xs0
864 where
865 not_ordered _ [] = False
866 not_ordered x (y : _) = x >= y
867 {-# INLINE not_ordered #-}
868
869 fromList' t0 xs = foldlStrict ins t0 xs
870 where ins t x = insert x t
871
872 go !_ t [] = t
873 go _ t [x] = insertMax x t
874 go s l xs@(x : xss) | not_ordered x xss = fromList' l xs
875 | otherwise = case create s xss of
876 (r, ys, []) -> go (s `shiftL` 1) (link x l r) ys
877 (r, _, ys) -> fromList' (link x l r) ys
878
879 -- The create is returning a triple (tree, xs, ys). Both xs and ys
880 -- represent not yet processed elements and only one of them can be nonempty.
881 -- If ys is nonempty, the keys in ys are not ordered with respect to tree
882 -- and must be inserted using fromList'. Otherwise the keys have been
883 -- ordered so far.
884 create !_ [] = (Tip, [], [])
885 create s xs@(x : xss)
886 | s == 1 = if not_ordered x xss then (Bin 1 x Tip Tip, [], xss)
887 else (Bin 1 x Tip Tip, xss, [])
888 | otherwise = case create (s `shiftR` 1) xs of
889 res@(_, [], _) -> res
890 (l, [y], zs) -> (insertMax y l, [], zs)
891 (l, ys@(y:yss), _) | not_ordered y yss -> (l, [], ys)
892 | otherwise -> case create (s `shiftR` 1) yss of
893 (r, zs, ws) -> (link y l r, zs, ws)
894 #if __GLASGOW_HASKELL__
895 {-# INLINABLE fromList #-}
896 #endif
897
898 {--------------------------------------------------------------------
899 Building trees from ascending/descending lists can be done in linear time.
900
901 Note that if [xs] is ascending that:
902 fromAscList xs == fromList xs
903 --------------------------------------------------------------------}
904 -- | /O(n)/. Build a set from an ascending list in linear time.
905 -- /The precondition (input list is ascending) is not checked./
906 fromAscList :: Eq a => [a] -> Set a
907 fromAscList xs
908 = fromDistinctAscList (combineEq xs)
909 where
910 -- [combineEq xs] combines equal elements with [const] in an ordered list [xs]
911 combineEq xs'
912 = case xs' of
913 [] -> []
914 [x] -> [x]
915 (x:xx) -> combineEq' x xx
916
917 combineEq' z [] = [z]
918 combineEq' z (x:xs')
919 | z==x = combineEq' z xs'
920 | otherwise = z:combineEq' x xs'
921 #if __GLASGOW_HASKELL__
922 {-# INLINABLE fromAscList #-}
923 #endif
924
925
926 -- | /O(n)/. Build a set from an ascending list of distinct elements in linear time.
927 -- /The precondition (input list is strictly ascending) is not checked./
928
929 -- For some reason, when 'singleton' is used in fromDistinctAscList or in
930 -- create, it is not inlined, so we inline it manually.
931 fromDistinctAscList :: [a] -> Set a
932 fromDistinctAscList [] = Tip
933 fromDistinctAscList (x0 : xs0) = go (1::Int) (Bin 1 x0 Tip Tip) xs0
934 where
935 go !_ t [] = t
936 go s l (x : xs) = case create s xs of
937 (r, ys) -> go (s `shiftL` 1) (link x l r) ys
938
939 create !_ [] = (Tip, [])
940 create s xs@(x : xs')
941 | s == 1 = (Bin 1 x Tip Tip, xs')
942 | otherwise = case create (s `shiftR` 1) xs of
943 res@(_, []) -> res
944 (l, y:ys) -> case create (s `shiftR` 1) ys of
945 (r, zs) -> (link y l r, zs)
946
947 {--------------------------------------------------------------------
948 Eq converts the set to a list. In a lazy setting, this
949 actually seems one of the faster methods to compare two trees
950 and it is certainly the simplest :-)
951 --------------------------------------------------------------------}
952 instance Eq a => Eq (Set a) where
953 t1 == t2 = (size t1 == size t2) && (toAscList t1 == toAscList t2)
954
955 {--------------------------------------------------------------------
956 Ord
957 --------------------------------------------------------------------}
958
959 instance Ord a => Ord (Set a) where
960 compare s1 s2 = compare (toAscList s1) (toAscList s2)
961
962 {--------------------------------------------------------------------
963 Show
964 --------------------------------------------------------------------}
965 instance Show a => Show (Set a) where
966 showsPrec p xs = showParen (p > 10) $
967 showString "fromList " . shows (toList xs)
968
969 {--------------------------------------------------------------------
970 Read
971 --------------------------------------------------------------------}
972 instance (Read a, Ord a) => Read (Set a) where
973 #ifdef __GLASGOW_HASKELL__
974 readPrec = parens $ prec 10 $ do
975 Ident "fromList" <- lexP
976 xs <- readPrec
977 return (fromList xs)
978
979 readListPrec = readListPrecDefault
980 #else
981 readsPrec p = readParen (p > 10) $ \ r -> do
982 ("fromList",s) <- lex r
983 (xs,t) <- reads s
984 return (fromList xs,t)
985 #endif
986
987 {--------------------------------------------------------------------
988 Typeable/Data
989 --------------------------------------------------------------------}
990
991 INSTANCE_TYPEABLE1(Set)
992
993 {--------------------------------------------------------------------
994 NFData
995 --------------------------------------------------------------------}
996
997 instance NFData a => NFData (Set a) where
998 rnf Tip = ()
999 rnf (Bin _ y l r) = rnf y `seq` rnf l `seq` rnf r
1000
1001 {--------------------------------------------------------------------
1002 Utility functions that return sub-ranges of the original
1003 tree. Some functions take a `Maybe value` as an argument to
1004 allow comparisons against infinite values. These are called `blow`
1005 (Nothing is -\infty) and `bhigh` (here Nothing is +\infty).
1006 We use MaybeS value, which is a Maybe strict in the Just case.
1007
1008 [trim blow bhigh t] A tree that is either empty or where [x > blow]
1009 and [x < bhigh] for the value [x] of the root.
1010 [filterGt blow t] A tree where for all values [k]. [k > blow]
1011 [filterLt bhigh t] A tree where for all values [k]. [k < bhigh]
1012
1013 [split k t] Returns two trees [l] and [r] where all values
1014 in [l] are <[k] and all keys in [r] are >[k].
1015 [splitMember k t] Just like [split] but also returns whether [k]
1016 was found in the tree.
1017 --------------------------------------------------------------------}
1018
1019 data MaybeS a = NothingS | JustS !a
1020
1021 {--------------------------------------------------------------------
1022 [trim blo bhi t] trims away all subtrees that surely contain no
1023 values between the range [blo] to [bhi]. The returned tree is either
1024 empty or the key of the root is between @blo@ and @bhi@.
1025 --------------------------------------------------------------------}
1026 trim :: Ord a => MaybeS a -> MaybeS a -> Set a -> Set a
1027 trim NothingS NothingS t = t
1028 trim (JustS lx) NothingS t = greater lx t where greater lo (Bin _ x _ r) | x <= lo = greater lo r
1029 greater _ t' = t'
1030 trim NothingS (JustS hx) t = lesser hx t where lesser hi (Bin _ x l _) | x >= hi = lesser hi l
1031 lesser _ t' = t'
1032 trim (JustS lx) (JustS hx) t = middle lx hx t where middle lo hi (Bin _ x _ r) | x <= lo = middle lo hi r
1033 middle lo hi (Bin _ x l _) | x >= hi = middle lo hi l
1034 middle _ _ t' = t'
1035 #if __GLASGOW_HASKELL__
1036 {-# INLINABLE trim #-}
1037 #endif
1038
1039 {--------------------------------------------------------------------
1040 [filterGt b t] filter all values >[b] from tree [t]
1041 [filterLt b t] filter all values <[b] from tree [t]
1042 --------------------------------------------------------------------}
1043 filterGt :: Ord a => MaybeS a -> Set a -> Set a
1044 filterGt NothingS t = t
1045 filterGt (JustS b) t = filter' b t
1046 where filter' _ Tip = Tip
1047 filter' b' (Bin _ x l r) =
1048 case compare b' x of LT -> link x (filter' b' l) r
1049 EQ -> r
1050 GT -> filter' b' r
1051 #if __GLASGOW_HASKELL__
1052 {-# INLINABLE filterGt #-}
1053 #endif
1054
1055 filterLt :: Ord a => MaybeS a -> Set a -> Set a
1056 filterLt NothingS t = t
1057 filterLt (JustS b) t = filter' b t
1058 where filter' _ Tip = Tip
1059 filter' b' (Bin _ x l r) =
1060 case compare x b' of LT -> link x l (filter' b' r)
1061 EQ -> l
1062 GT -> filter' b' l
1063 #if __GLASGOW_HASKELL__
1064 {-# INLINABLE filterLt #-}
1065 #endif
1066
1067 {--------------------------------------------------------------------
1068 Split
1069 --------------------------------------------------------------------}
1070 -- | /O(log n)/. The expression (@'split' x set@) is a pair @(set1,set2)@
1071 -- where @set1@ comprises the elements of @set@ less than @x@ and @set2@
1072 -- comprises the elements of @set@ greater than @x@.
1073 split :: Ord a => a -> Set a -> (Set a,Set a)
1074 split x0 t0 = toPair $ go x0 t0
1075 where
1076 go _ Tip = (Tip :*: Tip)
1077 go x (Bin _ y l r)
1078 = case compare x y of
1079 LT -> let (lt :*: gt) = go x l in (lt :*: link y gt r)
1080 GT -> let (lt :*: gt) = go x r in (link y l lt :*: gt)
1081 EQ -> (l :*: r)
1082 #if __GLASGOW_HASKELL__
1083 {-# INLINABLE split #-}
1084 #endif
1085
1086 -- | /O(log n)/. Performs a 'split' but also returns whether the pivot
1087 -- element was found in the original set.
1088 splitMember :: Ord a => a -> Set a -> (Set a,Bool,Set a)
1089 splitMember _ Tip = (Tip, False, Tip)
1090 splitMember x (Bin _ y l r)
1091 = case compare x y of
1092 LT -> let (lt, found, gt) = splitMember x l
1093 !gt' = link y gt r
1094 in (lt, found, gt')
1095 GT -> let (lt, found, gt) = splitMember x r
1096 !lt' = link y l lt
1097 in (lt', found, gt)
1098 EQ -> (l, True, r)
1099 #if __GLASGOW_HASKELL__
1100 {-# INLINABLE splitMember #-}
1101 #endif
1102
1103 {--------------------------------------------------------------------
1104 Indexing
1105 --------------------------------------------------------------------}
1106
1107 -- | /O(log n)/. Return the /index/ of an element, which is its zero-based
1108 -- index in the sorted sequence of elements. The index is a number from /0/ up
1109 -- to, but not including, the 'size' of the set. Calls 'error' when the element
1110 -- is not a 'member' of the set.
1111 --
1112 -- > findIndex 2 (fromList [5,3]) Error: element is not in the set
1113 -- > findIndex 3 (fromList [5,3]) == 0
1114 -- > findIndex 5 (fromList [5,3]) == 1
1115 -- > findIndex 6 (fromList [5,3]) Error: element is not in the set
1116
1117 -- See Note: Type of local 'go' function
1118 findIndex :: Ord a => a -> Set a -> Int
1119 findIndex = go 0
1120 where
1121 go :: Ord a => Int -> a -> Set a -> Int
1122 go !_ !_ Tip = error "Set.findIndex: element is not in the set"
1123 go idx x (Bin _ kx l r) = case compare x kx of
1124 LT -> go idx x l
1125 GT -> go (idx + size l + 1) x r
1126 EQ -> idx + size l
1127 #if __GLASGOW_HASKELL__
1128 {-# INLINABLE findIndex #-}
1129 #endif
1130
1131 -- | /O(log n)/. Lookup the /index/ of an element, which is its zero-based index in
1132 -- the sorted sequence of elements. The index is a number from /0/ up to, but not
1133 -- including, the 'size' of the set.
1134 --
1135 -- > isJust (lookupIndex 2 (fromList [5,3])) == False
1136 -- > fromJust (lookupIndex 3 (fromList [5,3])) == 0
1137 -- > fromJust (lookupIndex 5 (fromList [5,3])) == 1
1138 -- > isJust (lookupIndex 6 (fromList [5,3])) == False
1139
1140 -- See Note: Type of local 'go' function
1141 lookupIndex :: Ord a => a -> Set a -> Maybe Int
1142 lookupIndex = go 0
1143 where
1144 go :: Ord a => Int -> a -> Set a -> Maybe Int
1145 go !_ !_ Tip = Nothing
1146 go idx x (Bin _ kx l r) = case compare x kx of
1147 LT -> go idx x l
1148 GT -> go (idx + size l + 1) x r
1149 EQ -> Just $! idx + size l
1150 #if __GLASGOW_HASKELL__
1151 {-# INLINABLE lookupIndex #-}
1152 #endif
1153
1154 -- | /O(log n)/. Retrieve an element by its /index/, i.e. by its zero-based
1155 -- index in the sorted sequence of elements. If the /index/ is out of range (less
1156 -- than zero, greater or equal to 'size' of the set), 'error' is called.
1157 --
1158 -- > elemAt 0 (fromList [5,3]) == 3
1159 -- > elemAt 1 (fromList [5,3]) == 5
1160 -- > elemAt 2 (fromList [5,3]) Error: index out of range
1161
1162 elemAt :: Int -> Set a -> a
1163 elemAt !_ Tip = error "Set.elemAt: index out of range"
1164 elemAt i (Bin _ x l r)
1165 = case compare i sizeL of
1166 LT -> elemAt i l
1167 GT -> elemAt (i-sizeL-1) r
1168 EQ -> x
1169 where
1170 sizeL = size l
1171
1172 -- | /O(log n)/. Delete the element at /index/, i.e. by its zero-based index in
1173 -- the sorted sequence of elements. If the /index/ is out of range (less than zero,
1174 -- greater or equal to 'size' of the set), 'error' is called.
1175 --
1176 -- > deleteAt 0 (fromList [5,3]) == singleton 5
1177 -- > deleteAt 1 (fromList [5,3]) == singleton 3
1178 -- > deleteAt 2 (fromList [5,3]) Error: index out of range
1179 -- > deleteAt (-1) (fromList [5,3]) Error: index out of range
1180
1181 deleteAt :: Int -> Set a -> Set a
1182 deleteAt !i t =
1183 case t of
1184 Tip -> error "Set.deleteAt: index out of range"
1185 Bin _ x l r -> case compare i sizeL of
1186 LT -> balanceR x (deleteAt i l) r
1187 GT -> balanceL x l (deleteAt (i-sizeL-1) r)
1188 EQ -> glue l r
1189 where
1190 sizeL = size l
1191
1192
1193 {--------------------------------------------------------------------
1194 Utility functions that maintain the balance properties of the tree.
1195 All constructors assume that all values in [l] < [x] and all values
1196 in [r] > [x], and that [l] and [r] are valid trees.
1197
1198 In order of sophistication:
1199 [Bin sz x l r] The type constructor.
1200 [bin x l r] Maintains the correct size, assumes that both [l]
1201 and [r] are balanced with respect to each other.
1202 [balance x l r] Restores the balance and size.
1203 Assumes that the original tree was balanced and
1204 that [l] or [r] has changed by at most one element.
1205 [link x l r] Restores balance and size.
1206
1207 Furthermore, we can construct a new tree from two trees. Both operations
1208 assume that all values in [l] < all values in [r] and that [l] and [r]
1209 are valid:
1210 [glue l r] Glues [l] and [r] together. Assumes that [l] and
1211 [r] are already balanced with respect to each other.
1212 [merge l r] Merges two trees and restores balance.
1213 --------------------------------------------------------------------}
1214
1215 {--------------------------------------------------------------------
1216 Link
1217 --------------------------------------------------------------------}
1218 link :: a -> Set a -> Set a -> Set a
1219 link x Tip r = insertMin x r
1220 link x l Tip = insertMax x l
1221 link x l@(Bin sizeL y ly ry) r@(Bin sizeR z lz rz)
1222 | delta*sizeL < sizeR = balanceL z (link x l lz) rz
1223 | delta*sizeR < sizeL = balanceR y ly (link x ry r)
1224 | otherwise = bin x l r
1225
1226
1227 -- insertMin and insertMax don't perform potentially expensive comparisons.
1228 insertMax,insertMin :: a -> Set a -> Set a
1229 insertMax x t
1230 = case t of
1231 Tip -> singleton x
1232 Bin _ y l r
1233 -> balanceR y l (insertMax x r)
1234
1235 insertMin x t
1236 = case t of
1237 Tip -> singleton x
1238 Bin _ y l r
1239 -> balanceL y (insertMin x l) r
1240
1241 {--------------------------------------------------------------------
1242 [merge l r]: merges two trees.
1243 --------------------------------------------------------------------}
1244 merge :: Set a -> Set a -> Set a
1245 merge Tip r = r
1246 merge l Tip = l
1247 merge l@(Bin sizeL x lx rx) r@(Bin sizeR y ly ry)
1248 | delta*sizeL < sizeR = balanceL y (merge l ly) ry
1249 | delta*sizeR < sizeL = balanceR x lx (merge rx r)
1250 | otherwise = glue l r
1251
1252 {--------------------------------------------------------------------
1253 [glue l r]: glues two trees together.
1254 Assumes that [l] and [r] are already balanced with respect to each other.
1255 --------------------------------------------------------------------}
1256 glue :: Set a -> Set a -> Set a
1257 glue Tip r = r
1258 glue l Tip = l
1259 glue l r
1260 | size l > size r = let (m,l') = deleteFindMax l in balanceR m l' r
1261 | otherwise = let (m,r') = deleteFindMin r in balanceL m l r'
1262
1263 -- | /O(log n)/. Delete and find the minimal element.
1264 --
1265 -- > deleteFindMin set = (findMin set, deleteMin set)
1266
1267 deleteFindMin :: Set a -> (a,Set a)
1268 deleteFindMin t
1269 = case t of
1270 Bin _ x Tip r -> (x,r)
1271 Bin _ x l r -> let (xm,l') = deleteFindMin l in (xm,balanceR x l' r)
1272 Tip -> (error "Set.deleteFindMin: can not return the minimal element of an empty set", Tip)
1273
1274 -- | /O(log n)/. Delete and find the maximal element.
1275 --
1276 -- > deleteFindMax set = (findMax set, deleteMax set)
1277 deleteFindMax :: Set a -> (a,Set a)
1278 deleteFindMax t
1279 = case t of
1280 Bin _ x l Tip -> (x,l)
1281 Bin _ x l r -> let (xm,r') = deleteFindMax r in (xm,balanceL x l r')
1282 Tip -> (error "Set.deleteFindMax: can not return the maximal element of an empty set", Tip)
1283
1284 -- | /O(log n)/. Retrieves the minimal key of the set, and the set
1285 -- stripped of that element, or 'Nothing' if passed an empty set.
1286 minView :: Set a -> Maybe (a, Set a)
1287 minView Tip = Nothing
1288 minView x = Just (deleteFindMin x)
1289
1290 -- | /O(log n)/. Retrieves the maximal key of the set, and the set
1291 -- stripped of that element, or 'Nothing' if passed an empty set.
1292 maxView :: Set a -> Maybe (a, Set a)
1293 maxView Tip = Nothing
1294 maxView x = Just (deleteFindMax x)
1295
1296 {--------------------------------------------------------------------
1297 [balance x l r] balances two trees with value x.
1298 The sizes of the trees should balance after decreasing the
1299 size of one of them. (a rotation).
1300
1301 [delta] is the maximal relative difference between the sizes of
1302 two trees, it corresponds with the [w] in Adams' paper.
1303 [ratio] is the ratio between an outer and inner sibling of the
1304 heavier subtree in an unbalanced setting. It determines
1305 whether a double or single rotation should be performed
1306 to restore balance. It is correspondes with the inverse
1307 of $\alpha$ in Adam's article.
1308
1309 Note that according to the Adam's paper:
1310 - [delta] should be larger than 4.646 with a [ratio] of 2.
1311 - [delta] should be larger than 3.745 with a [ratio] of 1.534.
1312
1313 But the Adam's paper is errorneous:
1314 - it can be proved that for delta=2 and delta>=5 there does
1315 not exist any ratio that would work
1316 - delta=4.5 and ratio=2 does not work
1317
1318 That leaves two reasonable variants, delta=3 and delta=4,
1319 both with ratio=2.
1320
1321 - A lower [delta] leads to a more 'perfectly' balanced tree.
1322 - A higher [delta] performs less rebalancing.
1323
1324 In the benchmarks, delta=3 is faster on insert operations,
1325 and delta=4 has slightly better deletes. As the insert speedup
1326 is larger, we currently use delta=3.
1327
1328 --------------------------------------------------------------------}
1329 delta,ratio :: Int
1330 delta = 3
1331 ratio = 2
1332
1333 -- The balance function is equivalent to the following:
1334 --
1335 -- balance :: a -> Set a -> Set a -> Set a
1336 -- balance x l r
1337 -- | sizeL + sizeR <= 1 = Bin sizeX x l r
1338 -- | sizeR > delta*sizeL = rotateL x l r
1339 -- | sizeL > delta*sizeR = rotateR x l r
1340 -- | otherwise = Bin sizeX x l r
1341 -- where
1342 -- sizeL = size l
1343 -- sizeR = size r
1344 -- sizeX = sizeL + sizeR + 1
1345 --
1346 -- rotateL :: a -> Set a -> Set a -> Set a
1347 -- rotateL x l r@(Bin _ _ ly ry) | size ly < ratio*size ry = singleL x l r
1348 -- | otherwise = doubleL x l r
1349 -- rotateR :: a -> Set a -> Set a -> Set a
1350 -- rotateR x l@(Bin _ _ ly ry) r | size ry < ratio*size ly = singleR x l r
1351 -- | otherwise = doubleR x l r
1352 --
1353 -- singleL, singleR :: a -> Set a -> Set a -> Set a
1354 -- singleL x1 t1 (Bin _ x2 t2 t3) = bin x2 (bin x1 t1 t2) t3
1355 -- singleR x1 (Bin _ x2 t1 t2) t3 = bin x2 t1 (bin x1 t2 t3)
1356 --
1357 -- doubleL, doubleR :: a -> Set a -> Set a -> Set a
1358 -- doubleL x1 t1 (Bin _ x2 (Bin _ x3 t2 t3) t4) = bin x3 (bin x1 t1 t2) (bin x2 t3 t4)
1359 -- doubleR x1 (Bin _ x2 t1 (Bin _ x3 t2 t3)) t4 = bin x3 (bin x2 t1 t2) (bin x1 t3 t4)
1360 --
1361 -- It is only written in such a way that every node is pattern-matched only once.
1362 --
1363 -- Only balanceL and balanceR are needed at the moment, so balance is not here anymore.
1364 -- In case it is needed, it can be found in Data.Map.
1365
1366 -- Functions balanceL and balanceR are specialised versions of balance.
1367 -- balanceL only checks whether the left subtree is too big,
1368 -- balanceR only checks whether the right subtree is too big.
1369
1370 -- balanceL is called when left subtree might have been inserted to or when
1371 -- right subtree might have been deleted from.
1372 balanceL :: a -> Set a -> Set a -> Set a
1373 balanceL x l r = case r of
1374 Tip -> case l of
1375 Tip -> Bin 1 x Tip Tip
1376 (Bin _ _ Tip Tip) -> Bin 2 x l Tip
1377 (Bin _ lx Tip (Bin _ lrx _ _)) -> Bin 3 lrx (Bin 1 lx Tip Tip) (Bin 1 x Tip Tip)
1378 (Bin _ lx ll@(Bin _ _ _ _) Tip) -> Bin 3 lx ll (Bin 1 x Tip Tip)
1379 (Bin ls lx ll@(Bin lls _ _ _) lr@(Bin lrs lrx lrl lrr))
1380 | lrs < ratio*lls -> Bin (1+ls) lx ll (Bin (1+lrs) x lr Tip)
1381 | otherwise -> Bin (1+ls) lrx (Bin (1+lls+size lrl) lx ll lrl) (Bin (1+size lrr) x lrr Tip)
1382
1383 (Bin rs _ _ _) -> case l of
1384 Tip -> Bin (1+rs) x Tip r
1385
1386 (Bin ls lx ll lr)
1387 | ls > delta*rs -> case (ll, lr) of
1388 (Bin lls _ _ _, Bin lrs lrx lrl lrr)
1389 | lrs < ratio*lls -> Bin (1+ls+rs) lx ll (Bin (1+rs+lrs) x lr r)
1390 | otherwise -> Bin (1+ls+rs) lrx (Bin (1+lls+size lrl) lx ll lrl) (Bin (1+rs+size lrr) x lrr r)
1391 (_, _) -> error "Failure in Data.Map.balanceL"
1392 | otherwise -> Bin (1+ls+rs) x l r
1393 {-# NOINLINE balanceL #-}
1394
1395 -- balanceR is called when right subtree might have been inserted to or when
1396 -- left subtree might have been deleted from.
1397 balanceR :: a -> Set a -> Set a -> Set a
1398 balanceR x l r = case l of
1399 Tip -> case r of
1400 Tip -> Bin 1 x Tip Tip
1401 (Bin _ _ Tip Tip) -> Bin 2 x Tip r
1402 (Bin _ rx Tip rr@(Bin _ _ _ _)) -> Bin 3 rx (Bin 1 x Tip Tip) rr
1403 (Bin _ rx (Bin _ rlx _ _) Tip) -> Bin 3 rlx (Bin 1 x Tip Tip) (Bin 1 rx Tip Tip)
1404 (Bin rs rx rl@(Bin rls rlx rll rlr) rr@(Bin rrs _ _ _))
1405 | rls < ratio*rrs -> Bin (1+rs) rx (Bin (1+rls) x Tip rl) rr
1406 | otherwise -> Bin (1+rs) rlx (Bin (1+size rll) x Tip rll) (Bin (1+rrs+size rlr) rx rlr rr)
1407
1408 (Bin ls _ _ _) -> case r of
1409 Tip -> Bin (1+ls) x l Tip
1410
1411 (Bin rs rx rl rr)
1412 | rs > delta*ls -> case (rl, rr) of
1413 (Bin rls rlx rll rlr, Bin rrs _ _ _)
1414 | rls < ratio*rrs -> Bin (1+ls+rs) rx (Bin (1+ls+rls) x l rl) rr
1415 | otherwise -> Bin (1+ls+rs) rlx (Bin (1+ls+size rll) x l rll) (Bin (1+rrs+size rlr) rx rlr rr)
1416 (_, _) -> error "Failure in Data.Map.balanceR"
1417 | otherwise -> Bin (1+ls+rs) x l r
1418 {-# NOINLINE balanceR #-}
1419
1420 {--------------------------------------------------------------------
1421 The bin constructor maintains the size of the tree
1422 --------------------------------------------------------------------}
1423 bin :: a -> Set a -> Set a -> Set a
1424 bin x l r
1425 = Bin (size l + size r + 1) x l r
1426 {-# INLINE bin #-}
1427
1428
1429 {--------------------------------------------------------------------
1430 Utilities
1431 --------------------------------------------------------------------}
1432
1433 -- | /O(1)/. Decompose a set into pieces based on the structure of the underlying
1434 -- tree. This function is useful for consuming a set in parallel.
1435 --
1436 -- No guarantee is made as to the sizes of the pieces; an internal, but
1437 -- deterministic process determines this. However, it is guaranteed that the pieces
1438 -- returned will be in ascending order (all elements in the first subset less than all
1439 -- elements in the second, and so on).
1440 --
1441 -- Examples:
1442 --
1443 -- > splitRoot (fromList [1..6]) ==
1444 -- > [fromList [1,2,3],fromList [4],fromList [5,6]]
1445 --
1446 -- > splitRoot empty == []
1447 --
1448 -- Note that the current implementation does not return more than three subsets,
1449 -- but you should not depend on this behaviour because it can change in the
1450 -- future without notice.
1451 splitRoot :: Set a -> [Set a]
1452 splitRoot orig =
1453 case orig of
1454 Tip -> []
1455 Bin _ v l r -> [l, singleton v, r]
1456 {-# INLINE splitRoot #-}
1457
1458
1459 {--------------------------------------------------------------------
1460 Debugging
1461 --------------------------------------------------------------------}
1462 -- | /O(n)/. Show the tree that implements the set. The tree is shown
1463 -- in a compressed, hanging format.
1464 showTree :: Show a => Set a -> String
1465 showTree s
1466 = showTreeWith True False s
1467
1468
1469 {- | /O(n)/. The expression (@showTreeWith hang wide map@) shows
1470 the tree that implements the set. If @hang@ is
1471 @True@, a /hanging/ tree is shown otherwise a rotated tree is shown. If
1472 @wide@ is 'True', an extra wide version is shown.
1473
1474 > Set> putStrLn $ showTreeWith True False $ fromDistinctAscList [1..5]
1475 > 4
1476 > +--2
1477 > | +--1
1478 > | +--3
1479 > +--5
1480 >
1481 > Set> putStrLn $ showTreeWith True True $ fromDistinctAscList [1..5]
1482 > 4
1483 > |
1484 > +--2
1485 > | |
1486 > | +--1
1487 > | |
1488 > | +--3
1489 > |
1490 > +--5
1491 >
1492 > Set> putStrLn $ showTreeWith False True $ fromDistinctAscList [1..5]
1493 > +--5
1494 > |
1495 > 4
1496 > |
1497 > | +--3
1498 > | |
1499 > +--2
1500 > |
1501 > +--1
1502
1503 -}
1504 showTreeWith :: Show a => Bool -> Bool -> Set a -> String
1505 showTreeWith hang wide t
1506 | hang = (showsTreeHang wide [] t) ""
1507 | otherwise = (showsTree wide [] [] t) ""
1508
1509 showsTree :: Show a => Bool -> [String] -> [String] -> Set a -> ShowS
1510 showsTree wide lbars rbars t
1511 = case t of
1512 Tip -> showsBars lbars . showString "|\n"
1513 Bin _ x Tip Tip
1514 -> showsBars lbars . shows x . showString "\n"
1515 Bin _ x l r
1516 -> showsTree wide (withBar rbars) (withEmpty rbars) r .
1517 showWide wide rbars .
1518 showsBars lbars . shows x . showString "\n" .
1519 showWide wide lbars .
1520 showsTree wide (withEmpty lbars) (withBar lbars) l
1521
1522 showsTreeHang :: Show a => Bool -> [String] -> Set a -> ShowS
1523 showsTreeHang wide bars t
1524 = case t of
1525 Tip -> showsBars bars . showString "|\n"
1526 Bin _ x Tip Tip
1527 -> showsBars bars . shows x . showString "\n"
1528 Bin _ x l r
1529 -> showsBars bars . shows x . showString "\n" .
1530 showWide wide bars .
1531 showsTreeHang wide (withBar bars) l .
1532 showWide wide bars .
1533 showsTreeHang wide (withEmpty bars) r
1534
1535 showWide :: Bool -> [String] -> String -> String
1536 showWide wide bars
1537 | wide = showString (concat (reverse bars)) . showString "|\n"
1538 | otherwise = id
1539
1540 showsBars :: [String] -> ShowS
1541 showsBars bars
1542 = case bars of
1543 [] -> id
1544 _ -> showString (concat (reverse (tail bars))) . showString node
1545
1546 node :: String
1547 node = "+--"
1548
1549 withBar, withEmpty :: [String] -> [String]
1550 withBar bars = "| ":bars
1551 withEmpty bars = " ":bars
1552
1553 {--------------------------------------------------------------------
1554 Assertions
1555 --------------------------------------------------------------------}
1556 -- | /O(n)/. Test if the internal set structure is valid.
1557 valid :: Ord a => Set a -> Bool
1558 valid t
1559 = balanced t && ordered t && validsize t
1560
1561 ordered :: Ord a => Set a -> Bool
1562 ordered t
1563 = bounded (const True) (const True) t
1564 where
1565 bounded lo hi t'
1566 = case t' of
1567 Tip -> True
1568 Bin _ x l r -> (lo x) && (hi x) && bounded lo (<x) l && bounded (>x) hi r
1569
1570 balanced :: Set a -> Bool
1571 balanced t
1572 = case t of
1573 Tip -> True
1574 Bin _ _ l r -> (size l + size r <= 1 || (size l <= delta*size r && size r <= delta*size l)) &&
1575 balanced l && balanced r
1576
1577 validsize :: Set a -> Bool
1578 validsize t
1579 = (realsize t == Just (size t))
1580 where
1581 realsize t'
1582 = case t' of
1583 Tip -> Just 0
1584 Bin sz _ l r -> case (realsize l,realsize r) of
1585 (Just n,Just m) | n+m+1 == sz -> Just sz
1586 _ -> Nothing