Move nth method to RedBlack. Inline factories for tree nodes.

3 files changed, 20 insertions, 18 deletions

diff --git a/src/library/scala/collection/immutable/RedBlack.scala b/src/library/scala/collection/immutable/RedBlack.scala
index 6af6b6ef03..5729260cb2 100644
--- a/src/library/scala/collection/immutable/RedBlack.scala
+++ b/src/library/scala/collection/immutable/RedBlack.scala
@@ -79,6 +79,14 @@ object RedBlack {
   def keysIterator[A, _](tree: Tree[A, _]): Iterator[A] = new KeysIterator(tree)
   def valuesIterator[_, B](tree: Tree[_, B]): Iterator[B] = new ValuesIterator(tree)
 
+  @tailrec
+  def nth[A, B](tree: Tree[A, B], n: Int): Tree[A, B] = {
+    val count = RedBlack.count(tree.left)
+    if (n < count) nth(tree.left, n)
+    else if (n > count) nth(tree.right, n - count - 1)
+    else tree
+  }
+
   private[this] def balanceLeft[A, B, B1 >: B](isBlack: Boolean, z: A, zv: B, l: Tree[A, B1], d: Tree[A, B1]): Tree[A, B1] = {
     if (isRedTree(l) && isRedTree(l.left))
       RedTree(l.key, l.value, BlackTree(l.left.key, l.left.value, l.left.left, l.left.right), BlackTree(z, zv, l.right, d))
@@ -287,16 +295,9 @@ object RedBlack {
   extends Serializable {
     final val count: Int = 1 + RedBlack.count(left) + RedBlack.count(right)
     def isBlack: Boolean
-    def nth(n: Int): Tree[A, B] = {
-      val count = RedBlack.count(left)
-      if (n < count) left.nth(n)
-      else if (n > count) right.nth(n - count - 1)
-      else this
-    }
     def black: Tree[A, B]
     def red: Tree[A, B]
   }
-
   final class RedTree[A, +B](key: A,
                              value: B,
                              left: Tree[A, B],
@@ -306,10 +307,6 @@ object RedBlack {
     override def red = this
     override def toString = "RedTree(" + key + ", " + value + ", " + left + ", " + right + ")"
   }
-  object RedTree {
-    def apply[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = new RedTree(key, value, left, right)
-    def unapply[A, B](t: RedTree[A, B]) = Some((t.key, t.value, t.left, t.right))
-  }
   final class BlackTree[A, +B](key: A,
                                value: B,
                                left: Tree[A, B],
@@ -319,8 +316,13 @@ object RedBlack {
     override def red = RedTree(key, value, left, right)
     override def toString = "BlackTree(" + key + ", " + value + ", " + left + ", " + right + ")"
   }
+
+  object RedTree {
+    @inline def apply[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = new RedTree(key, value, left, right)
+    def unapply[A, B](t: RedTree[A, B]) = Some((t.key, t.value, t.left, t.right))
+  }
   object BlackTree {
-    def apply[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = new BlackTree(key, value, left, right)
+    @inline def apply[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = new BlackTree(key, value, left, right)
     def unapply[A, B](t: BlackTree[A, B]) = Some((t.key, t.value, t.left, t.right))
   }
 
diff --git a/src/library/scala/collection/immutable/TreeMap.scala b/src/library/scala/collection/immutable/TreeMap.scala
index 6e8cf625f4..7e22c19e11 100644
--- a/src/library/scala/collection/immutable/TreeMap.scala
+++ b/src/library/scala/collection/immutable/TreeMap.scala
@@ -88,20 +88,20 @@ class TreeMap[A, +B] private (tree: RedBlack.Tree[A, B])(implicit val ordering:
   override def drop(n: Int) = {
     if (n <= 0) this
     else if (n >= size) empty
-    else from(tree.nth(n).key)
+    else from(RB.nth(tree, n).key)
   }
 
   override def take(n: Int) = {
     if (n <= 0) empty
     else if (n >= size) this
-    else until(tree.nth(n).key)
+    else until(RB.nth(tree, n).key)
   }
 
   override def slice(from: Int, until: Int) = {
     if (until <= from) empty
     else if (from <= 0) take(until)
     else if (until >= size) drop(from)
-    else range(tree.nth(from).key, tree.nth(until).key)
+    else range(RB.nth(tree, from).key, RB.nth(tree, until).key)
   }
 
   override def dropRight(n: Int) = take(size - n)
diff --git a/src/library/scala/collection/immutable/TreeSet.scala b/src/library/scala/collection/immutable/TreeSet.scala
index 7c27e9f5b0..d36bc374c2 100644
--- a/src/library/scala/collection/immutable/TreeSet.scala
+++ b/src/library/scala/collection/immutable/TreeSet.scala
@@ -67,20 +67,20 @@ class TreeSet[A] private (tree: RedBlack.Tree[A, Unit])(implicit val ordering: O
   override def drop(n: Int) = {
     if (n <= 0) this
     else if (n >= size) empty
-    else from(tree.nth(n).key)
+    else from(RB.nth(tree, n).key)
   }
 
   override def take(n: Int) = {
     if (n <= 0) empty
     else if (n >= size) this
-    else until(tree.nth(n).key)
+    else until(RB.nth(tree, n).key)
   }
 
   override def slice(from: Int, until: Int) = {
     if (until <= from) empty
     else if (from <= 0) take(until)
     else if (until >= size) drop(from)
-    else range(tree.nth(from).key, tree.nth(until).key)
+    else range(RB.nth(tree, from).key, RB.nth(tree, until).key)
   }
 
   override def dropRight(n: Int) = take(size - n)