Skip to content

Commit 9388f29

Browse files
oderskyodersky
committed
Add mutable/TreeSet.scala to stdlib
1 parent 32f664b commit 9388f29

File tree

3 files changed

+251
-1
lines changed

3 files changed

+251
-1
lines changed

‎tests/pos-special/stdlib/collection/Set.scala

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,8 @@ trait Set[A]
2525
extends Iterable[A]
2626
with SetOps[A, Set, Set[A]]
2727
with Equals
28-
with IterableFactoryDefaults[A, Set] {
28+
with IterableFactoryDefaults[A, Set]
29+
with Pure {
2930
self: Set[A] =>
3031

3132
def canEqual(that: Any) = true

‎tests/pos-special/stdlib/collection/StrictOptimizedSetOps.scala

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
/*
2+
* Scala (https://www.scala-lang.org)
3+
*
4+
* Copyright EPFL and Lightbend, Inc.
5+
*
6+
* Licensed under Apache License 2.0
7+
* (http://www.apache.org/licenses/LICENSE-2.0).
8+
*
9+
* See the NOTICE file distributed with this work for
10+
* additional information regarding copyright ownership.
11+
*/
12+
13+
package scala.collection
14+
import language.experimental.captureChecking
15+
16+
/**
17+
* Trait that overrides set operations to take advantage of strict builders.
18+
*
19+
* @tparam A Elements type
20+
* @tparam CC Collection type constructor
21+
* @tparam C Collection type
22+
*/
23+
trait StrictOptimizedSetOps[A, +CC[_], +C <: SetOps[A, CC, C]]
24+
extends SetOps[A, CC, C]
25+
with StrictOptimizedIterableOps[A, CC, C] {
26+
27+
override def concat(that: IterableOnce[A]): C =
28+
strictOptimizedConcat(that, newSpecificBuilder)
29+
30+
}

‎tests/pos-special/stdlib/collection/mutable/TreeSet.scala

Lines changed: 219 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,219 @@
1+
/*
2+
* Scala (https://www.scala-lang.org)
3+
*
4+
* Copyright EPFL and Lightbend, Inc.
5+
*
6+
* Licensed under Apache License 2.0
7+
* (http://www.apache.org/licenses/LICENSE-2.0).
8+
*
9+
* See the NOTICE file distributed with this work for
10+
* additional information regarding copyright ownership.
11+
*/
12+
13+
package scala
14+
package collection.mutable
15+
16+
import scala.collection.Stepper.EfficientSplit
17+
import scala.collection.generic.DefaultSerializable
18+
import scala.collection.mutable.{RedBlackTree => RB}
19+
import scala.collection.{SortedIterableFactory, SortedSetFactoryDefaults, Stepper, StepperShape, StrictOptimizedIterableOps, StrictOptimizedSortedSetOps, mutable}
20+
import language.experimental.captureChecking
21+
22+
/**
23+
* A mutable sorted set implemented using a mutable red-black tree as underlying data structure.
24+
*
25+
* @param ordering the implicit ordering used to compare objects of type `A`.
26+
* @tparam A the type of the keys contained in this tree set.
27+
*
28+
* @define Coll mutable.TreeSet
29+
* @define coll mutable tree set
30+
*/
31+
// Original API designed in part by Lucien Pereira
32+
sealed class TreeSet[A] private (private val tree: RB.Tree[A, Null])(implicit val ordering: Ordering[A])
33+
extends AbstractSet[A]
34+
with SortedSet[A]
35+
with SortedSetOps[A, TreeSet, TreeSet[A]]
36+
with StrictOptimizedIterableOps[A, Set, TreeSet[A]]
37+
with StrictOptimizedSortedSetOps[A, TreeSet, TreeSet[A]]
38+
with SortedSetFactoryDefaults[A, TreeSet, Set]
39+
with DefaultSerializable {
40+
41+
if (ordering eq null)
42+
throw new NullPointerException("ordering must not be null")
43+
44+
/**
45+
* Creates an empty `TreeSet`.
46+
* @param ord the implicit ordering used to compare objects of type `A`.
47+
* @return an empty `TreeSet`.
48+
*/
49+
def this()(implicit ord: Ordering[A]) = this(RB.Tree.empty)(ord)
50+
51+
override def sortedIterableFactory: SortedIterableFactory[TreeSet] = TreeSet
52+
53+
def iterator: collection.Iterator[A] = RB.keysIterator(tree)
54+
55+
def iteratorFrom(start: A): collection.Iterator[A] = RB.keysIterator(tree, Some(start))
56+
57+
override def stepper[S <: Stepper[_]](implicit shape: StepperShape[A, S]): S with EfficientSplit = {
58+
import scala.collection.convert.impl._
59+
type T = RB.Node[A, Null]
60+
val s = shape.shape match {
61+
case StepperShape.IntShape => IntBinaryTreeStepper.from[T] (size, tree.root, _.left, _.right, _.key.asInstanceOf[Int])
62+
case StepperShape.LongShape => LongBinaryTreeStepper.from[T] (size, tree.root, _.left, _.right, _.key.asInstanceOf[Long])
63+
case StepperShape.DoubleShape => DoubleBinaryTreeStepper.from[T](size, tree.root, _.left, _.right, _.key.asInstanceOf[Double])
64+
case _ => shape.parUnbox(AnyBinaryTreeStepper.from[A, T](size, tree.root, _.left, _.right, _.key))
65+
}
66+
s.asInstanceOf[S with EfficientSplit]
67+
}
68+
69+
def addOne(elem: A): this.type = {
70+
RB.insert(tree, elem, null)
71+
this
72+
}
73+
74+
def subtractOne(elem: A): this.type = {
75+
RB.delete(tree, elem)
76+
this
77+
}
78+
79+
def clear(): Unit = RB.clear(tree)
80+
81+
def contains(elem: A): Boolean = RB.contains(tree, elem)
82+
83+
def unconstrained: collection.Set[A] = this
84+
85+
def rangeImpl(from: Option[A], until: Option[A]): TreeSet[A] = new TreeSetProjection(from, until)
86+
87+
override protected[this] def className: String = "TreeSet"
88+
89+
override def size: Int = RB.size(tree)
90+
override def knownSize: Int = size
91+
override def isEmpty: Boolean = RB.isEmpty(tree)
92+
93+
override def head: A = RB.minKey(tree).get
94+
95+
override def last: A = RB.maxKey(tree).get
96+
97+
override def minAfter(key: A): Option[A] = RB.minKeyAfter(tree, key)
98+
99+
override def maxBefore(key: A): Option[A] = RB.maxKeyBefore(tree, key)
100+
101+
override def foreach[U](f: A => U): Unit = RB.foreachKey(tree, f)
102+
103+
104+
/**
105+
* A ranged projection of a [[TreeSet]]. Mutations on this set affect the original set and vice versa.
106+
*
107+
* Only keys between this projection's key range will ever appear as elements of this set, independently of whether
108+
* the elements are added through the original set or through this view. That means that if one inserts an element in
109+
* a view whose key is outside the view's bounds, calls to `contains` will _not_ consider the newly added element.
110+
* Mutations are always reflected in the original set, though.
111+
*
112+
* @param from the lower bound (inclusive) of this projection wrapped in a `Some`, or `None` if there is no lower
113+
* bound.
114+
* @param until the upper bound (exclusive) of this projection wrapped in a `Some`, or `None` if there is no upper
115+
* bound.
116+
*/
117+
private[this] final class TreeSetProjection(from: Option[A], until: Option[A]) extends TreeSet[A](tree) {
118+
119+
/**
120+
* Given a possible new lower bound, chooses and returns the most constraining one (the maximum).
121+
*/
122+
private[this] def pickLowerBound(newFrom: Option[A]): Option[A] = (from, newFrom) match {
123+
case (Some(fr), Some(newFr)) => Some(ordering.max(fr, newFr))
124+
case (None, _) => newFrom
125+
case _ => from
126+
}
127+
128+
/**
129+
* Given a possible new upper bound, chooses and returns the most constraining one (the minimum).
130+
*/
131+
private[this] def pickUpperBound(newUntil: Option[A]): Option[A] = (until, newUntil) match {
132+
case (Some(unt), Some(newUnt)) => Some(ordering.min(unt, newUnt))
133+
case (None, _) => newUntil
134+
case _ => until
135+
}
136+
137+
/**
138+
* Returns true if the argument is inside the view bounds (between `from` and `until`).
139+
*/
140+
private[this] def isInsideViewBounds(key: A): Boolean = {
141+
val afterFrom = from.isEmpty || ordering.compare(from.get, key) <= 0
142+
val beforeUntil = until.isEmpty || ordering.compare(key, until.get) < 0
143+
afterFrom && beforeUntil
144+
}
145+
146+
override def rangeImpl(from: Option[A], until: Option[A]): TreeSet[A] =
147+
new TreeSetProjection(pickLowerBound(from), pickUpperBound(until))
148+
149+
override def contains(key: A) = isInsideViewBounds(key) && RB.contains(tree, key)
150+
151+
override def iterator = RB.keysIterator(tree, from, until)
152+
override def iteratorFrom(start: A) = RB.keysIterator(tree, pickLowerBound(Some(start)), until)
153+
154+
override def size = if (RB.size(tree) == 0) 0 else iterator.length
155+
override def knownSize: Int = if (RB.size(tree) == 0) 0 else -1
156+
override def isEmpty: Boolean = RB.size(tree) == 0 || !iterator.hasNext
157+
158+
override def head: A = headOption.get
159+
override def headOption: Option[A] = {
160+
val elem = if (from.isDefined) RB.minKeyAfter(tree, from.get) else RB.minKey(tree)
161+
(elem, until) match {
162+
case (Some(e), Some(unt)) if ordering.compare(e, unt) >= 0 => None
163+
case _ => elem
164+
}
165+
}
166+
167+
override def last: A = lastOption.get
168+
override def lastOption = {
169+
val elem = if (until.isDefined) RB.maxKeyBefore(tree, until.get) else RB.maxKey(tree)
170+
(elem, from) match {
171+
case (Some(e), Some(fr)) if ordering.compare(e, fr) < 0 => None
172+
case _ => elem
173+
}
174+
}
175+
176+
// Using the iterator should be efficient enough; if performance is deemed a problem later, a specialized
177+
// `foreachKey(f, from, until)` method can be created in `RedBlackTree`. See
178+
// https://github.com/scala/scala/pull/4608#discussion_r34307985 for a discussion about this.
179+
override def foreach[U](f: A => U): Unit = iterator.foreach(f)
180+
181+
override def clone(): mutable.TreeSet[A] = super.clone().rangeImpl(from, until)
182+
183+
}
184+
185+
}
186+
187+
/**
188+
* $factoryInfo
189+
* @define Coll `mutable.TreeSet`
190+
* @define coll mutable tree set
191+
*/
192+
@SerialVersionUID(3L)
193+
object TreeSet extends SortedIterableFactory[TreeSet] {
194+
195+
def empty[A : Ordering]: TreeSet[A] = new TreeSet[A]()
196+
197+
def from[E](it: IterableOnce[E]^)(implicit ordering: Ordering[E]): TreeSet[E] =
198+
it match {
199+
case ts: TreeSet[E] if ordering == ts.ordering =>
200+
new TreeSet[E](ts.tree.treeCopy())
201+
case ss: scala.collection.SortedSet[E] if ordering == ss.ordering =>
202+
new TreeSet[E](RB.fromOrderedKeys(ss.iterator, ss.size))
203+
case r: Range if (ordering eq Ordering.Int) || (ordering eq Ordering.Int.reverse) =>
204+
val it = if((ordering eq Ordering.Int) == (r.step > 0)) r.iterator else r.reverseIterator
205+
new TreeSet[E](RB.fromOrderedKeys(it.asInstanceOf[Iterator[E]], r.size))
206+
case _ =>
207+
val t: RB.Tree[E, Null] = RB.Tree.empty
208+
val i = it.iterator
209+
while (i.hasNext) RB.insert(t, i.next(), null)
210+
new TreeSet[E](t)
211+
}
212+
213+
def newBuilder[A](implicit ordering: Ordering[A]): Builder[A, TreeSet[A]] = new ReusableBuilder[A, TreeSet[A]] {
214+
private[this] var tree: RB.Tree[A, Null] = RB.Tree.empty
215+
def addOne(elem: A): this.type = { RB.insert(tree, elem, null); this }
216+
def result(): TreeSet[A] = new TreeSet[A](tree)
217+
def clear(): Unit = { tree = RB.Tree.empty }
218+
}
219+
}

0 commit comments

Comments
 (0)