Merge pull request #1204 from chuvoks/rb

Rank based take/drop/slice

2 files changed, 59 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]
   }
diff --git a/test/files/scalacheck/redblacktree.scala b/test/files/scalacheck/redblacktree.scala
index e2609fa200..bc7f92aa1b 100644
--- a/test/files/scalacheck/redblacktree.scala
+++ b/test/files/scalacheck/redblacktree.scala
@@ -205,6 +205,45 @@ package scala.collection.immutable.redblacktree {
       filteredTree == keysIterator(newTree).toList
     }
   }
+  
+  object TestDrop extends RedBlackTreeTest with RedBlackTreeInvariants  {
+    import RB._
+    
+    override type ModifyParm = Int
+    override def genParm(tree: Tree[String, Int]): Gen[ModifyParm] = choose(0, iterator(tree).size)
+    override def modify(tree: Tree[String, Int], parm: ModifyParm): Tree[String, Int] = drop(tree, parm)
+    
+    property("drop") = forAll(genInput) { case (tree, parm, newTree) =>
+      iterator(tree).drop(parm).toList == iterator(newTree).toList
+    }
+  }
+  
+  object TestTake extends RedBlackTreeTest with RedBlackTreeInvariants  {
+    import RB._
+    
+    override type ModifyParm = Int
+    override def genParm(tree: Tree[String, Int]): Gen[ModifyParm] = choose(0, iterator(tree).size)
+    override def modify(tree: Tree[String, Int], parm: ModifyParm): Tree[String, Int] = take(tree, parm)
+
+    property("take") = forAll(genInput) { case (tree, parm, newTree) =>
+      iterator(tree).take(parm).toList == iterator(newTree).toList
+    }
+  }
+  
+  object TestSlice extends RedBlackTreeTest with RedBlackTreeInvariants  {
+    import RB._
+
+    override type ModifyParm = (Int, Int)
+    override def genParm(tree: Tree[String, Int]): Gen[ModifyParm] = for {
+      from <- choose(0, iterator(tree).size)
+      to <- choose(from, iterator(tree).size)
+    } yield (from, to)
+    override def modify(tree: Tree[String, Int], parm: ModifyParm): Tree[String, Int] = slice(tree, parm._1, parm._2)
+    
+    property("slice") = forAll(genInput) { case (tree, parm, newTree) =>
+      iterator(tree).slice(parm._1, parm._2).toList == iterator(newTree).toList
+    }
+  }
 }
 
 object Test extends Properties("RedBlackTree") {
@@ -213,4 +252,7 @@ object Test extends Properties("RedBlackTree") {
   include(TestModify)
   include(TestDelete)
   include(TestRange)
+  include(TestDrop)
+  include(TestTake)
+  include(TestSlice)
 }