1 files changed, 53 insertions, 60 deletions

diff --git a/src/library/scala/collection/mutable/TreeSet.scala b/src/library/scala/collection/mutable/TreeSet.scala
index 5197af1b04..f849eea569 100644
--- a/src/library/scala/collection/mutable/TreeSet.scala
+++ b/src/library/scala/collection/mutable/TreeSet.scala
@@ -6,10 +6,13 @@
 **                          |/                                          **
 \*                                                                      */
 
-package scala.collection
+package scala
+package collection
 package mutable
 
 import generic._
+import scala.collection.immutable.{RedBlackTree => RB}
+import scala.runtime.ObjectRef
 
 /**
  * @define Coll `mutable.TreeSet`
@@ -29,95 +32,85 @@ object TreeSet extends MutableSortedSetFactory[TreeSet] {
 }
 
 /**
- * A mutable SortedSet using an immutable AVL Tree as underlying data structure.
+ * A mutable SortedSet using an immutable RedBlack Tree as underlying data structure.
  *
  * @author Lucien Pereira
  *
  */
-class TreeSet[A](implicit val ordering: Ordering[A]) extends SortedSet[A] with SetLike[A, TreeSet[A]]
+@deprecatedInheritance("TreeSet is not designed to enable meaningful subclassing.", "2.11.0")
+class TreeSet[A] private (treeRef: ObjectRef[RB.Tree[A, Null]], from: Option[A], until: Option[A])(implicit val ordering: Ordering[A])
+  extends SortedSet[A] with SetLike[A, TreeSet[A]]
   with SortedSetLike[A, TreeSet[A]] with Set[A] with Serializable {
 
-  // Projection constructor
-  private def this(base: Option[TreeSet[A]], from: Option[A], until: Option[A])(implicit ordering: Ordering[A]) {
-    this();
-    this.base = base
-    this.from = from
-    this.until = until
-  }
-
-  private var base: Option[TreeSet[A]] = None
-
-  private var from: Option[A] = None
-
-  private var until: Option[A] = None
-
-  private var avl: AVLTree[A] = Leaf
-
-  private var cardinality: Int = 0
-
-  def resolve: TreeSet[A] = base.getOrElse(this)
+  if (ordering eq null)
+    throw new NullPointerException("ordering must not be null")
 
-  private def isLeftAcceptable(from: Option[A], ordering: Ordering[A])(a: A): Boolean =
-    from.map(x => ordering.gteq(a, x)).getOrElse(true)
+  def this()(implicit ordering: Ordering[A]) = this(new ObjectRef(null), None, None)
 
-  private def isRightAcceptable(until: Option[A], ordering: Ordering[A])(a: A): Boolean =
-    until.map(x => ordering.lt(a, x)).getOrElse(true)
-
-  /**
-   * Cardinality store the set size, unfortunately a
-   * set view (given by rangeImpl)
-   * cannot take advantage of this optimisation
-   *
-   */
-  override def size: Int = base.map(_ => super.size).getOrElse(cardinality)
+  override def size: Int = RB.countInRange(treeRef.elem, from, until)
 
   override def stringPrefix = "TreeSet"
 
   override def empty: TreeSet[A] = TreeSet.empty
 
-  override def rangeImpl(from: Option[A], until: Option[A]): TreeSet[A] = new TreeSet(Some(this), from, until)
+  private def pickBound(comparison: (A, A) => A, oldBound: Option[A], newBound: Option[A]) = (newBound, oldBound) match {
+    case (Some(newB), Some(oldB)) => Some(comparison(newB, oldB))
+    case (None, _) => oldBound
+    case _ => newBound
+  }
+
+  override def rangeImpl(fromArg: Option[A], untilArg: Option[A]): TreeSet[A] = {
+    val newFrom = pickBound(ordering.max, fromArg, from)
+    val newUntil = pickBound(ordering.min, untilArg, until)
+
+    new TreeSet(treeRef, newFrom, newUntil)
+  }
 
   override def -=(elem: A): this.type = {
-    try {
-      resolve.avl = resolve.avl.remove(elem, ordering)
-      resolve.cardinality = resolve.cardinality - 1
-    } catch {
-      case e: NoSuchElementException => ()
-    }
+    treeRef.elem = RB.delete(treeRef.elem, elem)
     this
   }
 
   override def +=(elem: A): this.type = {
-    try {
-      resolve.avl = resolve.avl.insert(elem, ordering)
-      resolve.cardinality = resolve.cardinality + 1
-    } catch {
-      case e: IllegalArgumentException => ()
-    }
+    treeRef.elem = RB.update(treeRef.elem, elem, null, overwrite = false)
     this
   }
 
   /**
    * Thanks to the immutable nature of the
-   * underlying AVL Tree, we can share it with
+   * underlying Tree, we can share it with
    * the clone. So clone complexity in time is O(1).
    *
    */
-  override def clone(): TreeSet[A] = {
-    val clone = new TreeSet[A](base, from, until)
-    clone.avl = resolve.avl
-    clone.cardinality = resolve.cardinality
-    clone
-  }
+  override def clone(): TreeSet[A] =
+    new TreeSet[A](new ObjectRef(treeRef.elem), from, until)
+
+  private val notProjection = !(from.isDefined || until.isDefined)
 
   override def contains(elem: A): Boolean = {
-    isLeftAcceptable(from, ordering)(elem) &&
-    isRightAcceptable(until, ordering)(elem) &&
-    resolve.avl.contains(elem, ordering)
+    def leftAcceptable: Boolean = from match {
+      case Some(lb) => ordering.gteq(elem, lb)
+      case _ => true
+    }
+
+    def rightAcceptable: Boolean = until match {
+      case Some(ub) => ordering.lt(elem, ub)
+      case _ => true
+    }
+
+    (notProjection || (leftAcceptable && rightAcceptable)) &&
+      RB.contains(treeRef.elem, elem)
   }
 
-  override def iterator: Iterator[A] = resolve.avl.iterator
-    .dropWhile(e => !isLeftAcceptable(from, ordering)(e))
-      .takeWhile(e => isRightAcceptable(until, ordering)(e))
+  override def iterator: Iterator[A] = iteratorFrom(None)
 
+  override def keysIteratorFrom(start: A) = iteratorFrom(Some(start))
+
+  private def iteratorFrom(start: Option[A]) = {
+    val it = RB.keysIterator(treeRef.elem, pickBound(ordering.max, from, start))
+    until match {
+      case None => it
+      case Some(ub) => it takeWhile (k => ordering.lt(k, ub))
+    }
+  }
 }