From 78374f340e71d8e8f71c5bcd11452b72c207068c Mon Sep 17 00:00:00 2001
From: Erik Rozendaal <erik@deler.org>
Date: Sun, 22 Jan 2012 21:17:29 +0100
Subject: Custom implementations of drop/take/slice.

This mainly helps performance when comparing keys is expensive.
---
 .../scala/collection/immutable/RedBlackTree.scala  | 39 +++++++++++++++++++++-
 .../scala/collection/immutable/TreeMap.scala       |  6 ++--
 .../scala/collection/immutable/TreeSet.scala       |  6 ++--
 3 files changed, 44 insertions(+), 7 deletions(-)

(limited to 'src')

diff --git a/src/library/scala/collection/immutable/RedBlackTree.scala b/src/library/scala/collection/immutable/RedBlackTree.scala
index 7110ca4194..731a0f7975 100644
--- a/src/library/scala/collection/immutable/RedBlackTree.scala
+++ b/src/library/scala/collection/immutable/RedBlackTree.scala
@@ -56,6 +56,10 @@ object RedBlackTree {
   def to[A: Ordering, B](tree: Tree[A, B], to: A): Tree[A, B] = blacken(doTo(tree, to))
   def until[A: Ordering, B](tree: Tree[A, B], key: A): Tree[A, B] = blacken(doUntil(tree, key))
 
+  def drop[A: Ordering, B](tree: Tree[A, B], n: Int): Tree[A, B] = blacken(doDrop(tree, n))
+  def take[A: Ordering, B](tree: Tree[A, B], n: Int): Tree[A, B] = blacken(doTake(tree, n))
+  def slice[A: Ordering, B](tree: Tree[A, B], from: Int, until: Int): Tree[A, B] = blacken(doSlice(tree, from, until))
+
   def smallest[A, B](tree: Tree[A, B]): Tree[A, B] = {
     if (tree eq null) throw new NoSuchElementException("empty map")
     var result = tree
@@ -86,7 +90,7 @@ object RedBlackTree {
 
   @tailrec
   def nth[A, B](tree: Tree[A, B], n: Int): Tree[A, B] = {
-    val count = RedBlackTree.count(tree.left)
+    val count = this.count(tree.left)
     if (n < count) nth(tree.left, n)
     else if (n > count) nth(tree.right, n - count - 1)
     else tree
@@ -243,6 +247,39 @@ object RedBlackTree {
     else rebalance(tree, newLeft, newRight)
   }
 
+  private[this] def doDrop[A: Ordering, 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)
+    else rebalance(tree, newLeft, tree.right)
+  }
+  private[this] def doTake[A: Ordering, 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)
+    else rebalance(tree, tree.left, newRight)
+  }
+  private[this] def doSlice[A: Ordering, 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)
+    if (until <= count) return doSlice(tree.left, from, until)
+    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)
+    else if (newRight eq null) upd(newLeft, tree.key, tree.value)
+    else rebalance(tree, newLeft, newRight)
+  }
+
   // The zipper returned might have been traversed left-most (always the left child)
   // or right-most (always the right child). Left trees are traversed right-most,
   // and right trees are traversed leftmost.
diff --git a/src/library/scala/collection/immutable/TreeMap.scala b/src/library/scala/collection/immutable/TreeMap.scala
index a24221decc..dc4f79be35 100644
--- a/src/library/scala/collection/immutable/TreeMap.scala
+++ b/src/library/scala/collection/immutable/TreeMap.scala
@@ -89,20 +89,20 @@ class TreeMap[A, +B] private (tree: RB.Tree[A, B])(implicit val ordering: Orderi
   override def drop(n: Int) = {
     if (n <= 0) this
     else if (n >= size) empty
-    else from(RB.nth(tree, n).key)
+    else new TreeMap(RB.drop(tree, n))
   }
 
   override def take(n: Int) = {
     if (n <= 0) empty
     else if (n >= size) this
-    else until(RB.nth(tree, n).key)
+    else new TreeMap(RB.take(tree, n))
   }
 
   override def slice(from: Int, until: Int) = {
     if (until <= from) empty
     else if (from <= 0) take(until)
     else if (until >= size) drop(from)
-    else range(RB.nth(tree, from).key, RB.nth(tree, until).key)
+    else new TreeMap(RB.slice(tree, from, until))
   }
 
   override def dropRight(n: Int) = take(size - n)
diff --git a/src/library/scala/collection/immutable/TreeSet.scala b/src/library/scala/collection/immutable/TreeSet.scala
index e21aec362c..1b3d72ceb7 100644
--- a/src/library/scala/collection/immutable/TreeSet.scala
+++ b/src/library/scala/collection/immutable/TreeSet.scala
@@ -66,20 +66,20 @@ class TreeSet[A] private (tree: RB.Tree[A, Unit])(implicit val ordering: Orderin
   override def drop(n: Int) = {
     if (n <= 0) this
     else if (n >= size) empty
-    else from(RB.nth(tree, n).key)
+    else newSet(RB.drop(tree, n))
   }
 
   override def take(n: Int) = {
     if (n <= 0) empty
     else if (n >= size) this
-    else until(RB.nth(tree, n).key)
+    else newSet(RB.take(tree, n))
   }
 
   override def slice(from: Int, until: Int) = {
     if (until <= from) empty
     else if (from <= 0) take(until)
     else if (until >= size) drop(from)
-    else range(RB.nth(tree, from).key, RB.nth(tree, until).key)
+    else newSet(RB.slice(tree, from, until))
   }
 
   override def dropRight(n: Int) = take(size - n)
-- 
cgit v1.2.3