scala · paulp · Feb 16, 2012 · Dec 17, 2011 · Dec 17, 2011 · Dec 17, 2011
diff --git a/src/library/scala/collection/immutable/RedBlack.scala b/src/library/scala/collection/immutable/RedBlack.scala
@@ -11,10 +11,13 @@
 package scala.collection
 package immutable
 
-/** A base class containing the implementations for `TreeMaps` and `TreeSets`.
+/** Old base class that was used by previous implementations of `TreeMaps` and `TreeSets`.
+ *
+ *  Deprecated due to various performance bugs (see [[https://issues.scala-lang.org/browse/SI-5331 SI-5331]] for more information).
  *
  *  @since 2.3
  */
+@deprecated("use `TreeMap` or `TreeSet` instead", "2.10")
 @SerialVersionUID(8691885935445612921L)
 abstract class RedBlack[A] extends Serializable {
 
@@ -287,5 +290,3 @@ abstract class RedBlack[A] extends Serializable {
     def isBlack = true
   }
 }
-
-
diff --git a/src/library/scala/collection/immutable/RedBlackTree.scala b/src/library/scala/collection/immutable/RedBlackTree.scala
diff --git a/src/library/scala/collection/immutable/TreeMap.scala b/src/library/scala/collection/immutable/TreeMap.scala
@@ -12,6 +12,7 @@ package scala.collection
 package immutable
 
 import generic._
+import immutable.{RedBlackTree => RB}
 import mutable.Builder
 import annotation.bridge
 
@@ -23,7 +24,6 @@ object TreeMap extends ImmutableSortedMapFactory[TreeMap] {
   def empty[A, B](implicit ord: Ordering[A]) = new TreeMap[A, B]()(ord)
   /** $sortedMapCanBuildFromInfo */
   implicit def canBuildFrom[A, B](implicit ord: Ordering[A]): CanBuildFrom[Coll, (A, B), TreeMap[A, B]] = new SortedMapCanBuildFrom[A, B]
-  private def make[A, B](s: Int, t: RedBlack[A]#Tree[B])(implicit ord: Ordering[A]) = new TreeMap[A, B](s, t)(ord)
 }
 
 /** This class implements immutable maps using a tree.
@@ -46,31 +46,79 @@ object TreeMap extends ImmutableSortedMapFactory[TreeMap] {
  *  @define mayNotTerminateInf
  *  @define willNotTerminateInf
  */
-class TreeMap[A, +B](override val size: Int, t: RedBlack[A]#Tree[B])(implicit val ordering: Ordering[A])
-  extends RedBlack[A]
-     with SortedMap[A, B]
+class TreeMap[A, +B] private (tree: RB.Tree[A, B])(implicit val ordering: Ordering[A])
+  extends SortedMap[A, B]
      with SortedMapLike[A, B, TreeMap[A, B]]
      with MapLike[A, B, TreeMap[A, B]]
      with Serializable {
 
+  @deprecated("use `ordering.lt` instead", "2.10")
   def isSmaller(x: A, y: A) = ordering.lt(x, y)
 
   override protected[this] def newBuilder : Builder[(A, B), TreeMap[A, B]] =
     TreeMap.newBuilder[A, B]
 
-  def this()(implicit ordering: Ordering[A]) = this(0, null)(ordering)
+  override def size = RB.count(tree)
 
-  protected val tree: RedBlack[A]#Tree[B] = if (size == 0) Empty else t
+  def this()(implicit ordering: Ordering[A]) = this(null)(ordering)
 
-  override def rangeImpl(from : Option[A], until : Option[A]): TreeMap[A,B] = {
-    val ntree = tree.range(from,until)
-    new TreeMap[A,B](ntree.count, ntree)
-  }
+  override def rangeImpl(from: Option[A], until: Option[A]): TreeMap[A, B] = new TreeMap[A, B](RB.rangeImpl(tree, from, until))
+  override def range(from: A, until: A): TreeMap[A, B] = new TreeMap[A, B](RB.range(tree, from, until))
+  override def from(from: A): TreeMap[A, B] = new TreeMap[A, B](RB.from(tree, from))
+  override def to(to: A): TreeMap[A, B] = new TreeMap[A, B](RB.to(tree, to))
+  override def until(until: A): TreeMap[A, B] = new TreeMap[A, B](RB.until(tree, until))
 
-  override def firstKey = t.first
-  override def lastKey = t.last
+  override def firstKey = RB.smallest(tree).key
+  override def lastKey = RB.greatest(tree).key
   override def compare(k0: A, k1: A): Int = ordering.compare(k0, k1)
 
+  override def head = {
+    val smallest = RB.smallest(tree)
+    (smallest.key, smallest.value)
+  }
+  override def headOption = if (RB.isEmpty(tree)) None else Some(head)
+  override def last = {
+    val greatest = RB.greatest(tree)
+    (greatest.key, greatest.value)
+  }
+  override def lastOption = if (RB.isEmpty(tree)) None else Some(last)
+
+  override def tail = new TreeMap(RB.delete(tree, firstKey))
+  override def init = new TreeMap(RB.delete(tree, lastKey))
+
+  override def drop(n: Int) = {
+    if (n <= 0) this
+    else if (n >= size) empty
+    else new TreeMap(RB.drop(tree, n))
+  }
+
+  override def take(n: Int) = {
+    if (n <= 0) empty
+    else if (n >= size) this
+    else new TreeMap(RB.take(tree, n))
+  }
+
+  override def slice(from: Int, until: Int) = {
+    if (until <= from) empty
+    else if (from <= 0) take(until)
+    else if (until >= size) drop(from)
+    else new TreeMap(RB.slice(tree, from, until))
+  }
+
+  override def dropRight(n: Int) = take(size - n)
+  override def takeRight(n: Int) = drop(size - n)
+  override def splitAt(n: Int) = (take(n), drop(n))
+
+  private[this] def countWhile(p: ((A, B)) => Boolean): Int = {
+    var result = 0
+    val it = iterator
+    while (it.hasNext && p(it.next)) result += 1
+    result
+  }
+  override def dropWhile(p: ((A, B)) => Boolean) = drop(countWhile(p))
+  override def takeWhile(p: ((A, B)) => Boolean) = take(countWhile(p))
+  override def span(p: ((A, B)) => Boolean) = splitAt(countWhile(p))
+
   /** A factory to create empty maps of the same type of keys.
    */
   override def empty: TreeMap[A, B] = TreeMap.empty[A, B](ordering)
@@ -84,10 +132,7 @@ class TreeMap[A, +B](override val size: Int, t: RedBlack[A]#Tree[B])(implicit va
    *  @param value   the value to be associated with `key`
    *  @return        a new $coll with the updated binding
    */
-  override def updated [B1 >: B](key: A, value: B1): TreeMap[A, B1] = {
-    val newsize = if (tree.lookup(key).isEmpty) size + 1 else size
-    TreeMap.make(newsize, tree.update(key, value))
-  }
+  override def updated [B1 >: B](key: A, value: B1): TreeMap[A, B1] = new TreeMap(RB.update(tree, key, value))
 
   /** Add a key/value pair to this map.
    *  @tparam   B1   type of the value of the new binding, a supertype of `B`
@@ -128,36 +173,36 @@ class TreeMap[A, +B](override val size: Int, t: RedBlack[A]#Tree[B])(implicit va
    *  @return       a new $coll with the inserted binding, if it wasn't present in the map
    */
   def insert [B1 >: B](key: A, value: B1): TreeMap[A, B1] = {
-    assert(tree.lookup(key).isEmpty)
-    TreeMap.make(size + 1, tree.update(key, value))
+    assert(!RB.contains(tree, key))
+    new TreeMap(RB.update(tree, key, value))
   }
 
   def - (key:A): TreeMap[A, B] =
-    if (tree.lookup(key).isEmpty) this
-    else if (size == 1) empty
-    else TreeMap.make(size - 1, tree.delete(key))
+    if (!RB.contains(tree, key)) this
+    else new TreeMap(RB.delete(tree, key))
 
   /** Check if this map maps `key` to a value and return the
    *  value if it exists.
    *
    *  @param  key     the key of the mapping of interest
    *  @return         the value of the mapping, if it exists
    */
-  override def get(key: A): Option[B] = tree.lookup(key) match {
-    case n: NonEmpty[b] => Some(n.value)
-    case _ => None
-  }
+  override def get(key: A): Option[B] = RB.get(tree, key)
 
   /** Creates a new iterator over all elements contained in this
    *  object.
    *
    *  @return the new iterator
    */
-  def iterator: Iterator[(A, B)] = tree.toStream.iterator
+  override def iterator: Iterator[(A, B)] = RB.iterator(tree)
+
+  override def keysIterator: Iterator[A] = RB.keysIterator(tree)
+  override def valuesIterator: Iterator[B] = RB.valuesIterator(tree)
 
-  override def toStream: Stream[(A, B)] = tree.toStream
+  override def contains(key: A): Boolean = RB.contains(tree, key)
+  override def isDefinedAt(key: A): Boolean = RB.contains(tree, key)
 
-  override def foreach[U](f : ((A,B)) =>  U) = tree foreach { case (x, y) => f(x, y) }
+  override def foreach[U](f : ((A,B)) =>  U) = RB.foreach(tree, f)
 }
 
 

diff --git a/src/library/scala/collection/immutable/TreeSet.scala b/src/library/scala/collection/immutable/TreeSet.scala
@@ -12,6 +12,7 @@ package scala.collection
 package immutable
 
 import generic._
+import immutable.{RedBlackTree => RB}
 import mutable.{ Builder, SetBuilder }
 
 /** $factoryInfo
@@ -46,20 +47,61 @@ object TreeSet extends ImmutableSortedSetFactory[TreeSet] {
  *  @define mayNotTerminateInf
  *  @define willNotTerminateInf
  */
-@SerialVersionUID(-234066569443569402L)
-class TreeSet[A](override val size: Int, t: RedBlack[A]#Tree[Unit])
-                (implicit val ordering: Ordering[A])
-  extends RedBlack[A] with SortedSet[A] with SortedSetLike[A, TreeSet[A]] with Serializable {
+@SerialVersionUID(-5685982407650748405L)
+class TreeSet[A] private (tree: RB.Tree[A, Unit])(implicit val ordering: Ordering[A])
+  extends SortedSet[A] with SortedSetLike[A, TreeSet[A]] with Serializable {
 
   override def stringPrefix = "TreeSet"
 
-  
F438
def isSmaller(x: A, y: A) = compare(x,y) < 0
+  override def size = RB.count(tree)
+
+  override def head = RB.smallest(tree).key
+  override def headOption = if (RB.isEmpty(tree)) None else Some(head)
+  override def last = RB.greatest(tree).key
+  override def lastOption = if (RB.isEmpty(tree)) None else Some(last)
+
+  override def tail = new TreeSet(RB.delete(tree, firstKey))
+  override def init = new TreeSet(RB.delete(tree, lastKey))
+
+  override def drop(n: Int) = {
+    if (n <= 0) this
+    else if (n >= size) empty
+    else newSet(RB.drop(tree, n))
+  }
 
-  def this()(implicit ordering: Ordering[A]) = this(0, null)(ordering)
+  override def take(n: Int) = {
+    if (n <= 0) empty
+    else if (n >= size) this
+    else newSet(RB.take(tree, n))
+  }
 
-  protected val tree: RedBlack[A]#Tree[Unit] = if (size == 0) Empty else t
+  override def slice(from: Int, until: Int) = {
+    if (until <= from) empty
+    else if (from <= 0) take(until)
+    else if (until >= size) drop(from)
+    else newSet(RB.slice(tree, from, until))
+  }
 
-  private def newSet(s: Int, t: RedBlack[A]#Tree[Unit]) = new TreeSet[A](s, t)
+  override def dropRight(n: Int) = take(size - n)
+  override def takeRight(n: Int) = drop(size - n)
+  override def splitAt(n: Int) = (take(n), drop(n))
+
+  private[this] def countWhile(p: A => Boolean): Int = {
+    var result = 0
+    val it = iterator
+    while (it.hasNext && p(it.next)) result += 1
+    result
+  }
+  override def dropWhile(p: A => Boolean) = drop(countWhile(p))
+  override def takeWhile(p: A => Boolean) = take(countWhile(p))
+  override def span(p: A => Boolean) = splitAt(countWhile(p))
+
+  @deprecated("use `ordering.lt` instead", "2.10")
+  def isSmaller(x: A, y: A) = compare(x,y) < 0
+
+  def this()(implicit ordering: Ordering[A]) = this(null)(ordering)
+
+  private def newSet(t: RB.Tree[A, Unit]) = new TreeSet[A](t)
 
   /** A factory to create empty sets of the same type of keys.
    */
@@ -70,10 +112,7 @@ class TreeSet[A](override val size: Int, t: RedBlack[A]#Tree[Unit])
    *  @param elem    a new element to add.
    *  @return        a new $coll containing `elem` and all the elements of this $coll.
    */
-  def + (elem: A): TreeSet[A] = {
-    val newsize = if (tree.lookup(elem).isEmpty) size + 1 else size
-    newSet(newsize, tree.update(elem, ()))
-  }
+  def + (elem: A): TreeSet[A] = newSet(RB.update(tree, elem, ()))
 
   /** A new `TreeSet` with the entry added is returned,
    *  assuming that elem is <em>not</em> in the TreeSet.
@@ -82,8 +121,8 @@ class TreeSet[A](override val size: Int, t: RedBlack[A]#Tree[Unit])
    *  @return        a new $coll containing `elem` and all the elements of this $coll.
    */
   def insert(elem: A): TreeSet[A] = {
-    assert(tree.lookup(elem).isEmpty)
-    newSet(size + 1, tree.update(elem, ()))
+    assert(!RB.contains(tree, elem))
+    newSet(RB.update(tree, elem, ()))
   }
 
   /** Creates a new `TreeSet` with the entry removed.
@@ -92,31 +131,31 @@ class TreeSet[A](override val size: Int, t: RedBlack[A]#Tree[Unit])
    *  @return        a new $coll containing all the elements of this $coll except `elem`.
    */
   def - (elem:A): TreeSet[A] =
-    if (tree.lookup(elem).isEmpty) this
-    else newSet(size - 1, tree delete elem)
+    if (!RB.contains(tree, elem)) this
+    else newSet(RB.delete(tree, elem))
 
   /** Checks if this set contains element `elem`.
    *
    *  @param  elem    the element to check for membership.
    *  @return true, iff `elem` is contained in this set.
    */
-  def contains(elem: A): Boolean = !tree.lookup(elem).isEmpty
+  def contains(elem: A): Boolean = RB.contains(tree, elem)
 
   /** Creates a new iterator over all elements contained in this
    *  object.
    *
    *  @return the new iterator
    */
-  def iterator: Iterator[A] = tree.toStream.iterator map (_._1)
+  def iterator: Iterator[A] = RB.keysIterator(tree)
 
-  override def toStream: Stream[A] = tree.toStream map (_._1)
+  override def foreach[U](f: A =>  U) = RB.foreachKey(tree, f)
 
-  override def foreach[U](f: A =>  U) = tree foreach { (x, y) => f(x) }
+  override def rangeImpl(from: Option[A], until: Option[A]): TreeSet[A] = newSet(RB.rangeImpl(tree, from, until))
+  override def range(from: A, until: A): TreeSet[A] = newSet(RB.range(tree, from, until))
+  override def from(from: A): TreeSet[A] = newSet(RB.from(tree, from))
+  override def to(to: A): TreeSet[A] = newSet(RB.to(tree, to))
+  override def until(until: A): TreeSet[A] = newSet(RB.until(tree, until))
 
-  override def rangeImpl(from: Option[A], until: Option[A]): TreeSet[A] = {
-    val tree = this.tree.range(from, until)
-    newSet(tree.count, tree)
-  }
-  override def firstKey = tree.first
-  override def lastKey = tree.last
+  override def firstKey = head
+  override def lastKey = last
 }