SI-6642 Refactor mutable.TreeSet to use RedBlackTree instead of AVL

There was no reason to have mutable.TreeSet use AVLTree while
immutable.TreeSet and immutable.HashSet used RedBlackTree. In
particular that would have meant duplicating the iteratorFrom logic
unnecessarily. So this commit refactors mutable.TreeSet to use
RedBlackTree for everything, including iteratorFrom. It also adds
a test to make sure TreeSet works as expected.

AVLTree should be dead code since it's private[scala.collection.mutable]
and only used by mutable.TreeSet, but to be safe it's only deprecated
in this commit.

4 files changed, 223 insertions, 79 deletions

diff --git a/src/library/scala/collection/immutable/RedBlackTree.scala b/src/library/scala/collection/immutable/RedBlackTree.scala
index c0d46765be..d8c69f026b 100644
--- a/src/library/scala/collection/immutable/RedBlackTree.scala
+++ b/src/library/scala/collection/immutable/RedBlackTree.scala
@@ -24,7 +24,7 @@ import scala.annotation.meta.getter
  *
  *  @since 2.10
  */
-private[immutable]
+private[collection]
 object RedBlackTree {
 
   def isEmpty(tree: Tree[_, _]): Boolean = tree eq null
@@ -44,6 +44,25 @@ object RedBlackTree {
   }
 
   def count(tree: Tree[_, _]) = if (tree eq null) 0 else tree.count
+  /**
+   * Count all the nodes with keys greater than or equal to the lower bound and less than the upper bound.
+   * The two bounds are optional.
+   */
+  def countInRange[A, _](tree: Tree[A, _], from: Option[A], to:Option[A])(implicit ordering: Ordering[A]) : Int = 
+    if (tree eq null) 0 else
+    (from, to) match {
+      // with no bounds use this node's count
+      case (None, None) => tree.count
+      // if node is less than the lower bound, try the tree on the right, it might be in range
+      case (Some(lb), _) if ordering.lt(tree.key, lb) => countInRange(tree.right, from, to)
+      // if node is greater than or equal to the upper bound, try the tree on the left, it might be in range
+      case (_, Some(ub)) if ordering.gteq(tree.key, ub) => countInRange(tree.left, from, to)
+      // node is in range so the tree on the left will all be less than the upper bound and the tree on the
+      // right will all be greater than or equal to the lower bound. So 1 for this node plus
+      // count the subtrees by stripping off the bounds that we don't need any more
+      case _ => 1 + countInRange(tree.left, from, None) + countInRange(tree.right, None, to)
+    
+    }
   def update[A, B, B1 >: B](tree: Tree[A, B], k: A, v: B1, overwrite: Boolean)(implicit ordering: Ordering[A]): Tree[A, B1] = blacken(upd(tree, k, v, overwrite))
   def delete[A, B](tree: Tree[A, B], k: A)(implicit ordering: Ordering[A]): Tree[A, B] = blacken(del(tree, k))
   def rangeImpl[A: Ordering, B](tree: Tree[A, B], from: Option[A], until: Option[A]): Tree[A, B] = (from, until) match {
diff --git a/src/library/scala/collection/mutable/AVLTree.scala b/src/library/scala/collection/mutable/AVLTree.scala
index 157e5dae62..da63778fcc 100644
--- a/src/library/scala/collection/mutable/AVLTree.scala
+++ b/src/library/scala/collection/mutable/AVLTree.scala
@@ -15,7 +15,7 @@ package mutable
  * An immutable AVL Tree implementation used by mutable.TreeSet
  *
  * @author Lucien Pereira
- *
+ * @deprecated("AVLTree and its related classes are being removed from the standard library since they're not different enough from RedBlackTree to justify keeping them.", "2.11")
  */
 private[mutable] sealed trait AVLTree[+A] extends Serializable {
   def balance: Int
@@ -65,12 +65,18 @@ private[mutable] sealed trait AVLTree[+A] extends Serializable {
   def doubleRightRotation[B >: A]: Node[B] = sys.error("Should not happen.")
 }
 
+/**
+ * @deprecated("AVLTree and its related classes are being removed from the standard library since they're not different enough from RedBlackTree to justify keeping them.", "2.11")
+ */
 private case object Leaf extends AVLTree[Nothing] {
   override val balance: Int = 0
 
   override val depth: Int = -1
 }
 
+/**
+ * @deprecated("AVLTree and its related classes are being removed from the standard library since they're not different enough from RedBlackTree to justify keeping them.", "2.11")
+ */
 private case class Node[A](val data: A, val left: AVLTree[A], val right: AVLTree[A]) extends AVLTree[A] {
   override val balance: Int = right.depth - left.depth
 
@@ -205,6 +211,9 @@ private case class Node[A](val data: A, val left: AVLTree[A], val right: AVLTree
   }
 }
 
+/**
+ * @deprecated("AVLTree and its related classes are being removed from the standard library since they're not different enough from RedBlackTree to justify keeping them.", "2.11")
+ */
 private class AVLIterator[A](root: Node[A]) extends Iterator[A] {
   val stack = mutable.ArrayStack[Node[A]](root)
   diveLeft()
diff --git a/src/library/scala/collection/mutable/TreeSet.scala b/src/library/scala/collection/mutable/TreeSet.scala
index 4fd35658fa..9113d8221b 100644
--- a/src/library/scala/collection/mutable/TreeSet.scala
+++ b/src/library/scala/collection/mutable/TreeSet.scala
@@ -10,6 +10,8 @@ package scala.collection
 package mutable
 
 import generic._
+import scala.collection.immutable.{RedBlackTree => RB}
+import scala.runtime.ObjectRef
 
 /**
  * @define Coll `mutable.TreeSet`
@@ -29,112 +31,81 @@ 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]]
+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 this()(implicit ordering: Ordering[A]) = this(new ObjectRef(null), None, None)
 
-  def resolve: TreeSet[A] = base.getOrElse(this)
-
-  private def isLeftAcceptable(from: Option[A], ordering: Ordering[A])(a: A): Boolean =
-    from.map(x => ordering.gteq(a, x)).getOrElse(true)
-
-  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, 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)
   }
 
-  // TODO see the discussion on keysIteratorFrom
-  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)
   
-  // TODO because TreeSets are potentially ranged views into other TreeSets
-  // what this really needs to do is walk the whole stack of tree sets, find
-  // the highest "from", and then do a tree walk of the underlying avl tree
-  // to find that spot in max(O(stack depth), O(log tree.size)) time which
-  // should effectively be O(log size) since ranged views are rare and
-  // even more rarely deep. With the following implementation it's
-  // O(N log N) to get an iterator from a start point.  
-  // But before engaging that endeavor I think mutable.TreeSet should be
-  // based on the same immutable RedBlackTree that immutable.TreeSet is
-  // based on. There's no good reason to have these two collections based
-  // on two different balanced binary trees. That'll save
-  // having to duplicate logic for finding the starting point of a
-  // sorted binary tree iterator, logic that has already been
-  // baked into RedBlackTree.
-  override def keysIteratorFrom(start: A) = from(start).iterator
-
+  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))
+    }
+  }
 }
diff --git a/test/files/run/mutable-treeset.scala b/test/files/run/mutable-treeset.scala
new file mode 100644
index 0000000000..c9918cba96
--- /dev/null
+++ b/test/files/run/mutable-treeset.scala
@@ -0,0 +1,145 @@
+import scala.collection.mutable.TreeSet
+
+object Test extends App {
+  val list = List(6,5,4,3,2,1,1,2,3,4,5,6,6,5,4,3,2,1)
+  val distinct = list.distinct
+  val sorted = distinct.sorted
+
+  // sublist stuff for a single level of slicing
+  val min = list.min
+  val max = list.max
+  val nonlist = ((min - 10) until (max + 20) filterNot list.contains).toList
+  val sublist = list filter {x => x >=(min + 1) && x < max} 
+  val distinctSublist = sublist.distinct 
+  val subnonlist = min :: max :: nonlist
+  val subsorted = distinctSublist.sorted
+
+  // subsublist for a 2nd level of slicing
+  val almostmin = sublist.min
+  val almostmax = sublist.max
+  val subsublist = sublist filter {x => x >=(almostmin + 1) && x < almostmax} 
+  val distinctSubsublist = subsublist.distinct 
+  val subsubnonlist = almostmin :: almostmax :: subnonlist
+  val subsubsorted = distinctSubsublist.sorted
+
+  def testSize {
+    def check(set : TreeSet[Int], list: List[Int]) { 
+      assert(set.size == list.size, s"$set had size ${set.size} while $list had size ${list.size}")
+    }
+
+    check(TreeSet[Int](), List[Int]())
+    val set = TreeSet(list:_*)
+    check(set, distinct)
+    check(set.clone, distinct)
+
+    val subset = set from (min + 1) until max
+    check(subset, distinctSublist)
+    check(subset.clone, distinctSublist)
+
+    val subsubset = subset from (almostmin + 1) until almostmax
+    check(subsubset, distinctSubsublist)
+    check(subsubset.clone, distinctSubsublist)
+  }
+
+  def testContains {
+    def check(set : TreeSet[Int], list: List[Int], nonlist: List[Int]) {
+      assert(list forall set.apply, s"$set did not contain all elements of $list using apply")
+      assert(list forall set.contains, s"$set did not contain all elements of $list using contains")
+      assert(!(nonlist exists set.apply), s"$set had an element from $nonlist using apply")
+      assert(!(nonlist exists set.contains), s"$set had an element from $nonlist using contains")      
+    }
+
+    val set = TreeSet(list:_*)
+    check(set, list, nonlist)
+    check(set.clone, list, nonlist)
+
+    val subset = set from (min + 1) until max
+    check(subset, sublist, subnonlist)
+    check(subset.clone, sublist, subnonlist)
+
+    val subsubset = subset from (almostmin + 1) until almostmax
+    check(subsubset, subsublist, subsubnonlist)
+    check(subsubset.clone, subsublist, subsubnonlist)
+  }
+
+  def testAdd {
+    def check(set : TreeSet[Int], list: List[Int], nonlist: List[Int]) {
+      var builtList = List[Int]()
+      for (x <- list) {
+        set += x
+        builtList = (builtList :+ x).distinct.sorted filterNot nonlist.contains
+        assert(builtList forall set.apply, s"$set did not contain all elements of $builtList using apply")
+        assert(builtList.size == set.size, s"$set had size ${set.size} while $builtList had size ${builtList.size}")
+      }
+      assert(!(nonlist exists set.apply), s"$set had an element from $nonlist using apply")
+      assert(!(nonlist exists set.contains), s"$set had an element from $nonlist using contains")      
+    }
+
+    val set = TreeSet[Int]()
+    val clone = set.clone    
+    val subset = set.clone from (min + 1) until max
+    val subclone = subset.clone
+    val subsubset = subset.clone from (almostmin + 1) until almostmax
+    val subsubclone = subsubset.clone
+
+    check(set, list, nonlist)
+    check(clone, list, nonlist)
+
+    check(subset, list, subnonlist)
+    check(subclone, list, subnonlist)
+
+    check(subsubset, list, subsubnonlist)
+    check(subsubclone, list, subsubnonlist)
+  }
+
+  def testRemove {
+    def check(set: TreeSet[Int], sorted: List[Int]) {
+      var builtList = sorted
+      for (x <- list) {
+        set remove x
+        builtList = builtList filterNot (_ == x)
+        assert(builtList forall set.apply, s"$set did not contain all elements of $builtList using apply")
+        assert(builtList.size == set.size, s"$set had size $set.size while $builtList had size $builtList.size")
+      }   
+    }
+    val set = TreeSet(list:_*)
+    val clone = set.clone    
+    val subset = set.clone from (min + 1) until max
+    val subclone = subset.clone
+    val subsubset = subset.clone from (almostmin + 1) until almostmax
+    val subsubclone = subsubset.clone
+
+    check(set, sorted)
+    check(clone, sorted)
+
+    check(subset, subsorted)
+    check(subclone, subsorted)
+    
+    check(subsubset, subsubsorted)
+    check(subsubclone, subsubsorted)
+  }
+
+  def testIterator {
+    def check(set: TreeSet[Int], list: List[Int]) {
+      val it = set.iterator.toList
+      assert(it == list, s"$it did not equal $list")
+    }
+    val set = TreeSet(list: _*)
+    check(set, sorted)
+    check(set.clone, sorted)
+
+    val subset = set from (min + 1) until max
+    check(subset, subsorted)
+    check(subset.clone, subsorted)
+
+    val subsubset = subset from (almostmin + 1) until almostmax
+    check(subsubset, subsubsorted)
+    check(subsubset.clone, subsubsorted)
+  }
+
+  testSize
+  testContains
+  testAdd
+  testRemove
+  testIterator
+}