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