Tests for takeWhile/dropWhile/span.

Also simplified implementation of span to just use splitAt.

src/library/scala/collection/immutable/TreeMap.scala

@@ -116,10 +116,7 @@ class TreeMap[A, +B] private (tree: RB.Tree[A, B])(implicit val ordering: Orderi
   }
   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) = {
-    val n = countWhile(p)
-    (take(n), drop(n))
-  }
+  override def span(p: ((A, B)) => Boolean) = splitAt(countWhile(p))
 
   /** A factory to create empty maps of the same type of keys.
    */

src/library/scala/collection/immutable/TreeSet.scala

@@ -94,10 +94,7 @@ class TreeSet[A] private (tree: RB.Tree[A, Unit])(implicit val ordering: Orderin
   }
   override def dropWhile(p: A => Boolean) = drop(countWhile(p))
   override def takeWhile(p: A => Boolean) = take(countWhile(p))
-  override def span(p: A => Boolean) = {
-    val n = countWhile(p)
-    (take(n), drop(n))
-  }
+  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

test/files/scalacheck/treemap.scala

@@ -96,6 +96,21 @@ object Test extends Properties("TreeMap") {
     prefix == subject.take(n) && suffix == subject.drop(n)
   }
 
+  property("takeWhile") = forAll { (subject: TreeMap[Int, String]) =>
+    val result = subject.takeWhile(_._1 < 0)
+    result.forall(_._1 < 0) && result == subject.take(result.size)
+  }
+
+  property("dropWhile") = forAll { (subject: TreeMap[Int, String]) =>
+    val result = subject.dropWhile(_._1 < 0)
+    result.forall(_._1 >= 0) && result == subject.takeRight(result.size)
+  }
+
+  property("span identity") = forAll { (subject: TreeMap[Int, String]) =>
+    val (prefix, suffix) = subject.span(_._1 < 0)
+    prefix.forall(_._1 < 0) && suffix.forall(_._1 >= 0) && subject == prefix ++ suffix
+  }
+
   property("remove single") = forAll { (subject: TreeMap[Int, String]) => subject.nonEmpty ==> {
     val key = oneOf(subject.keys.toSeq).sample.get
     val removed = subject - key

test/files/scalacheck/treeset.scala

@@ -92,6 +92,21 @@ object Test extends Properties("TreeSet") {
     prefix == subject.take(n) && suffix == subject.drop(n)
   }
 
+  property("takeWhile") = forAll { (subject: TreeMap[Int, String]) =>
+    val result = subject.takeWhile(_._1 < 0)
+    result.forall(_._1 < 0) && result == subject.take(result.size)
+  }
+
+  property("dropWhile") = forAll { (subject: TreeMap[Int, String]) =>
+    val result = subject.dropWhile(_._1 < 0)
+    result.forall(_._1 >= 0) && result == subject.takeRight(result.size)
+  }
+
+  property("span identity") = forAll { (subject: TreeMap[Int, String]) =>
+    val (prefix, suffix) = subject.span(_._1 < 0)
+    prefix.forall(_._1 < 0) && suffix.forall(_._1 >= 0) && subject == prefix ++ suffix
+  }
+
   property("remove single") = forAll { (subject: TreeSet[Int]) => subject.nonEmpty ==> {
     val element = oneOf(subject.toSeq).sample.get
     val removed = subject - element