Custom coded version of range/from/to/until.

This avoids unnecessary allocation of Option and Function objects,
mostly helping performance of small trees.

3 files changed, 49 insertions, 26 deletions

diff --git a/src/library/scala/collection/immutable/RedBlackTree.scala b/src/library/scala/collection/immutable/RedBlackTree.scala
index d8caeab096..7110ca4194 100644
--- a/src/library/scala/collection/immutable/RedBlackTree.scala
+++ b/src/library/scala/collection/immutable/RedBlackTree.scala
@@ -45,11 +45,16 @@ object RedBlackTree {
   def count(tree: Tree[_, _]) = if (tree eq null) 0 else tree.count
   def update[A, B, B1 >: B](tree: Tree[A, B], k: A, v: B1)(implicit ordering: Ordering[A]): Tree[A, B1] = blacken(upd(tree, k, v))
   def delete[A, B](tree: Tree[A, B], k: A)(implicit ordering: Ordering[A]): Tree[A, B] = blacken(del(tree, k))
-  def range[A, B](tree: Tree[A, B], low: Option[A], lowInclusive: Boolean, high: Option[A], highInclusive: Boolean)(implicit ordering: Ordering[A]): Tree[A, B] = {
-    val after: Option[A => Boolean] = low.map(key => if (lowInclusive) ordering.lt(_, key) else ordering.lteq(_, key))
-    val before: Option[A => Boolean] = high.map(key => if (highInclusive) ordering.lt(key, _) else ordering.lteq(key, _))
-    blacken(rng(tree, after, before))
+  def rangeImpl[A: Ordering, B](tree: Tree[A, B], from: Option[A], until: Option[A]): Tree[A, B] = (from, until) match {
+    case (Some(from), Some(until)) => this.range(tree, from, until)
+    case (Some(from), None)        => this.from(tree, from)
+    case (None,       Some(until)) => this.until(tree, until)
+    case (None,       None)        => tree
   }
+  def range[A: Ordering, B](tree: Tree[A, B], from: A, until: A): Tree[A, B] = blacken(doRange(tree, from, until))
+  def from[A: Ordering, B](tree: Tree[A, B], from: A): Tree[A, B] = blacken(doFrom(tree, from))
+  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 smallest[A, B](tree: Tree[A, B]): Tree[A, B] = {
     if (tree eq null) throw new NoSuchElementException("empty map")
@@ -202,13 +207,36 @@ object RedBlackTree {
     else append(tree.left, tree.right)
   }
 
-  private[this] def rng[A, B](tree: Tree[A, B], after: Option[A => Boolean], before: Option[A => Boolean])(implicit ordering: Ordering[A]): Tree[A, B] = {
+  private[this] def doFrom[A, B](tree: Tree[A, B], from: A)(implicit ordering: Ordering[A]): Tree[A, B] = {
     if (tree eq null) return null
-    if (after == None && before == None) return tree
-    if (after != None && after.get(tree.key)) return rng(tree.right, after, before);
-    if (before != None && before.get(tree.key)) return rng(tree.left, after, before);
-    val newLeft = rng(tree.left, after, None)
-    val newRight = rng(tree.right, None, before)
+    if (ordering.lt(tree.key, from)) return doFrom(tree.right, from)
+    val newLeft = doFrom(tree.left, from)
+    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 doTo[A, B](tree: Tree[A, B], to: A)(implicit ordering: Ordering[A]): Tree[A, B] = {
+    if (tree eq null) return null
+    if (ordering.lt(to, tree.key)) return doTo(tree.left, to)
+    val newRight = doTo(tree.right, to)
+    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 doUntil[A, B](tree: Tree[A, B], until: A)(implicit ordering: Ordering[A]): Tree[A, B] = {
+    if (tree eq null) return null
+    if (ordering.lteq(until, tree.key)) return doUntil(tree.left, until)
+    val newRight = doUntil(tree.right, until)
+    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 doRange[A, B](tree: Tree[A, B], from: A, until: A)(implicit ordering: Ordering[A]): Tree[A, B] = {
+    if (tree eq null) return null
+    if (ordering.lt(tree.key, from)) return doRange(tree.right, from, until);
+    if (ordering.lteq(until, tree.key)) return doRange(tree.left, from, until);
+    val newLeft = doFrom(tree.left, from)
+    val newRight = doUntil(tree.right, until)
     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);
diff --git a/src/library/scala/collection/immutable/TreeMap.scala b/src/library/scala/collection/immutable/TreeMap.scala
index 3eba64dca3..a24221decc 100644
--- a/src/library/scala/collection/immutable/TreeMap.scala
+++ b/src/library/scala/collection/immutable/TreeMap.scala
@@ -62,14 +62,11 @@ class TreeMap[A, +B] private (tree: RB.Tree[A, B])(implicit val ordering: Orderi
 
   def this()(implicit ordering: Ordering[A]) = this(null)(ordering)
 
-  override def rangeImpl(from: Option[A], until: Option[A]): TreeMap[A, B] = {
-    val ntree = RB.range(tree, from, true, until, false)
-    new TreeMap[A, B](ntree)
-  }
-  override def to(to: A): TreeMap[A, B] = {
-    val ntree = RB.range(tree, None, true, Some(to), true)
-    new TreeMap[A, B](ntree)
-  }
+  override def rangeImpl(from: Option[A], until: Option[A]): TreeMap[A, B] = new TreeMap[A, B](RB.rangeImpl(tree, from, until))
+  override def range(from: A, until: A): TreeMap[A, B] = new TreeMap[A, B](RB.range(tree, from, until))
+  override def from(from: A): TreeMap[A, B] = new TreeMap[A, B](RB.from(tree, from))
+  override def to(to: A): TreeMap[A, B] = new TreeMap[A, B](RB.to(tree, to))
+  override def until(until: A): TreeMap[A, B] = new TreeMap[A, B](RB.until(tree, until))
 
   override def firstKey = RB.smallest(tree).key
   override def lastKey = RB.greatest(tree).key
diff --git a/src/library/scala/collection/immutable/TreeSet.scala b/src/library/scala/collection/immutable/TreeSet.scala
index 5dd80e87a4..e21aec362c 100644
--- a/src/library/scala/collection/immutable/TreeSet.scala
+++ b/src/library/scala/collection/immutable/TreeSet.scala
@@ -150,14 +150,12 @@ class TreeSet[A] private (tree: RB.Tree[A, Unit])(implicit val ordering: Orderin
 
   override def foreach[U](f: A =>  U) = RB.foreachKey(tree, f)
 
-  override def rangeImpl(from: Option[A], until: Option[A]): TreeSet[A] = {
-    val ntree = RB.range(tree, from, true, until, false)
-    newSet(ntree)
-  }
-  override def to(to: A): TreeSet[A] = {
-    val ntree = RB.range(tree, None, true, Some(to), true)
-    newSet(ntree)
-  }
+  override def rangeImpl(from: Option[A], until: Option[A]): TreeSet[A] = newSet(RB.rangeImpl(tree, from, until))
+  override def range(from: A, until: A): TreeSet[A] = newSet(RB.range(tree, from, until))
+  override def from(from: A): TreeSet[A] = newSet(RB.from(tree, from))
+  override def to(to: A): TreeSet[A] = newSet(RB.to(tree, to))
+  override def until(until: A): TreeSet[A] = newSet(RB.until(tree, until))
+
   override def firstKey = head
   override def lastKey = last
 }