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