Merge pull request #1204 from chuvoks/rb

Rank based take/drop/slice

1 files changed, 42 insertions, 0 deletions

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)
 }