From 8d678236d820619819e52f2497d1dd1df29f1184 Mon Sep 17 00:00:00 2001
From: Erik Rozendaal <erik@deler.org>
Date: Sun, 18 Dec 2011 22:49:39 +0100
Subject: Implemented takeWhile/dropWhile/span to use tree splitting. This
 changes the operation from O(n log n) to O(n) and allows for more structural
 sharing.

---
 src/library/scala/collection/immutable/TreeMap.scala | 13 +++++++++++++
 src/library/scala/collection/immutable/TreeSet.scala | 13 +++++++++++++
 2 files changed, 26 insertions(+)

(limited to 'src/library')

diff --git a/src/library/scala/collection/immutable/TreeMap.scala b/src/library/scala/collection/immutable/TreeMap.scala
index bc91bbe268..da2aef1c22 100644
--- a/src/library/scala/collection/immutable/TreeMap.scala
+++ b/src/library/scala/collection/immutable/TreeMap.scala
@@ -108,6 +108,19 @@ class TreeMap[A, +B](override val size: Int, t: RedBlack[A]#Tree[B])(implicit va
   override def takeRight(n: Int) = drop(size - n)
   override def splitAt(n: Int) = (take(n), drop(n))
 
+  private[this] def countWhile(p: ((A, B)) => Boolean): Int = {
+    var result = 0
+    val it = iterator
+    while (it.hasNext && p(it.next)) result += 1
+    result
+  }
+  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))
+  }
+
   /** A factory to create empty maps of the same type of keys.
    */
   override def empty: TreeMap[A, B] = TreeMap.empty[A, B](ordering)
diff --git a/src/library/scala/collection/immutable/TreeSet.scala b/src/library/scala/collection/immutable/TreeSet.scala
index dfaffcd581..7fb333959e 100644
--- a/src/library/scala/collection/immutable/TreeSet.scala
+++ b/src/library/scala/collection/immutable/TreeSet.scala
@@ -84,6 +84,19 @@ class TreeSet[A](override val size: Int, t: RedBlack[A]#Tree[Unit])
   override def takeRight(n: Int) = drop(size - n)
   override def splitAt(n: Int) = (take(n), drop(n))
 
+  private[this] def countWhile(p: A => Boolean): Int = {
+    var result = 0
+    val it = iterator
+    while (it.hasNext && p(it.next)) result += 1
+    result
+  }
+  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))
+  }
+
   def isSmaller(x: A, y: A) = compare(x,y) < 0
 
   def this()(implicit ordering: Ordering[A]) = this(0, null)(ordering)
-- 
cgit v1.2.3