Merge pull request #1204 from chuvoks/rb

Rank based take/drop/slice

1 files changed, 17 insertions, 8 deletions

diff --git a/src/library/scala/collection/immutable/RedBlackTree.scala b/src/library/scala/collection/immutable/RedBlackTree.scala
index 5bdba26d02..bb489dd80a 100644
--- a/src/library/scala/collection/immutable/RedBlackTree.scala
+++ b/src/library/scala/collection/immutable/RedBlackTree.scala
@@ -132,6 +132,15 @@ object RedBlackTree {
     else if (overwrite || k != tree.key) mkTree(isBlackTree(tree), k, v, tree.left, tree.right)
     else tree
   }
+  private[this] def updNth[A, B, B1 >: B](tree: Tree[A, B], idx: Int, k: A, v: B1, overwrite: Boolean): Tree[A, B1] = if (tree eq null) {
+    RedTree(k, v, null, null)
+  } else {
+    val rank = count(tree.left) + 1
+    if (idx < rank) balanceLeft(isBlackTree(tree), tree.key, tree.value, updNth(tree.left, idx, k, v, overwrite), tree.right)
+    else if (idx > rank) balanceRight(isBlackTree(tree), tree.key, tree.value, tree.left, updNth(tree.right, idx - rank, k, v, overwrite))
+    else if (overwrite) mkTree(isBlackTree(tree), k, v, tree.left, tree.right)
+    else tree
+  }
 
   /* Based on Stefan Kahrs' Haskell version of Okasaki's Red&Black Trees
    * http://www.cse.unsw.edu.au/~dons/data/RedBlackTree.html */
@@ -249,27 +258,27 @@ object RedBlackTree {
     else rebalance(tree, newLeft, newRight)
   }
 
-  private[this] def doDrop[A: Ordering, B](tree: Tree[A, B], n: Int): Tree[A, B] = {
+  private[this] def doDrop[A, B](tree: Tree[A, B], n: Int): Tree[A, B] = {
     if (n <= 0) return tree
     if (n >= this.count(tree)) return null
     val count = this.count(tree.left)
     if (n > count) return doDrop(tree.right, n - count - 1)
     val newLeft = doDrop(tree.left, n)
     if (newLeft eq tree.left) tree
-    else if (newLeft eq null) upd(tree.right, tree.key, tree.value, false)
+    else if (newLeft eq null) updNth(tree.right, n - count - 1, tree.key, tree.value, false)
     else rebalance(tree, newLeft, tree.right)
   }
-  private[this] def doTake[A: Ordering, B](tree: Tree[A, B], n: Int): Tree[A, B] = {
+  private[this] def doTake[A, B](tree: Tree[A, B], n: Int): Tree[A, B] = {
     if (n <= 0) return null
     if (n >= this.count(tree)) return tree
     val count = this.count(tree.left)
     if (n <= count) return doTake(tree.left, n)
     val newRight = doTake(tree.right, n - count - 1)
     if (newRight eq tree.right) tree
-    else if (newRight eq null) upd(tree.left, tree.key, tree.value, false)
+    else if (newRight eq null) updNth(tree.left, n, tree.key, tree.value, false)
     else rebalance(tree, tree.left, newRight)
   }
-  private[this] def doSlice[A: Ordering, B](tree: Tree[A, B], from: Int, until: Int): Tree[A, B] = {
+  private[this] def doSlice[A, B](tree: Tree[A, B], from: Int, until: Int): Tree[A, B] = {
     if (tree eq null) return null
     val count = this.count(tree.left)
     if (from > count) return doSlice(tree.right, from - count - 1, until - count - 1)
@@ -277,8 +286,8 @@ object RedBlackTree {
     val newLeft = doDrop(tree.left, from)
     val newRight = doTake(tree.right, until - count - 1)
     if ((newLeft eq tree.left) && (newRight eq tree.right)) tree
-    else if (newLeft eq null) upd(newRight, tree.key, tree.value, false)
-    else if (newRight eq null) upd(newLeft, tree.key, tree.value, false)
+    else if (newLeft eq null) updNth(newRight, from - count - 1, tree.key, tree.value, false)
+    else if (newRight eq null) updNth(newLeft, until, tree.key, tree.value, false)
     else rebalance(tree, newLeft, newRight)
   }
 
@@ -380,7 +389,7 @@ object RedBlackTree {
     @(inline @getter) final val left: Tree[A, B],
     @(inline @getter) final val right: Tree[A, B])
   extends Serializable {
-    final val count: Int = 1 + RedBlackTree.count(left) + RedBlackTree.count(right)
+    @(inline @getter) final val count: Int = 1 + RedBlackTree.count(left) + RedBlackTree.count(right)
     def black: Tree[A, B]
     def red: Tree[A, B]
   }