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

This avoids unnecessary allocation of Option and Function objects, mostly helping performance of small trees.
scala · paulp · Feb 16, 2012 · Dec 17, 2011 · Dec 17, 2011 · Dec 17, 2011
commit 7824dbd3cfe6704ab56aa5ceb2af2f5f4e55cbc7
diff --git 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
@@ -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
@@ -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
 }
diff --git a/test/files/scalacheck/redblacktree.scala b/test/files/scalacheck/redblacktree.scala
@@ -174,39 +174,33 @@ package scala.collection.immutable.redblacktree {
   object TestRange extends RedBlackTreeTest with RedBlackTreeInvariants  {
     import RB._
 
-    override type ModifyParm = (Option[Int], Boolean, Option[Int], Boolean)
+    override type ModifyParm = (Option[Int], Option[Int])
     override def genParm(tree: Tree[String, Int]): Gen[ModifyParm] = for {
       from <- choose(0, iterator(tree).size)
-      fromInclusive <- oneOf(false, true)
       to <- choose(0, iterator(tree).size) suchThat (from <=)
-      toInclusive <- oneOf(false, true)
       optionalFrom <- oneOf(Some(from), None, Some(from)) // Double Some(n) to get around a bug
       optionalTo <- oneOf(Some(to), None, Some(to)) // Double Some(n) to get around a bug
-    } yield (optionalFrom, fromInclusive, optionalTo, toInclusive)
+    } yield (optionalFrom, optionalTo)
 
     override def modify(tree: Tree[String, Int], parm: ModifyParm): Tree[String, Int] = {
       val from = parm._1 flatMap (nodeAt(tree, _) map (_._1))
-      val to = parm._3 flatMap (nodeAt(tree, _) map (_._1))
-      range(tree, from, parm._2, to, parm._4)
+      val to = parm._2 flatMap (nodeAt(tree, _) map (_._1))
+      rangeImpl(tree, from, to)
     }
 
     property("range boundaries respected") = forAll(genInput) { case (tree, parm, newTree) =>
       val from = parm._1 flatMap (nodeAt(tree, _) map (_._1))
-      val fromPredicate: String => String => Boolean = if (parm._2) (_ <=) else (_ <)
-      val to = parm._3 flatMap (nodeAt(tree, _) map (_._1))
-      val toPredicate: String => String => Boolean = if (parm._4) (_ >=) else (_ >)
-      ("lower boundary" |: (from forall ( key => keysIterator(newTree) forall fromPredicate(key)))) &&
-      ("upper boundary" |: (to forall ( key => keysIterator(newTree) forall toPredicate(key))))
+      val to = parm._2 flatMap (nodeAt(tree, _) map (_._1))
+      ("lower boundary" |: (from forall ( key => keysIterator(newTree) forall (key <=)))) &&
+      ("upper boundary" |: (to forall ( key => keysIterator(newTree) forall (key >))))
     }
 
     property("range returns all elements") = forAll(genInput) { case (tree, parm, newTree) =>
       val from = parm._1 flatMap (nodeAt(tree, _) map (_._1))
-      val fromPredicate: String => String => Boolean = if (parm._2) (_ >=) else (_ >)
-      val to = parm._3 flatMap (nodeAt(tree, _) map (_._1))
-      val toPredicate: String => String => Boolean = if (parm._4) (_ <=) else (_ <)
+      val to = parm._2 flatMap (nodeAt(tree, _) map (_._1))
       val filteredTree = (keysIterator(tree)
-        .filter(key => from forall fromPredicate(key))
-        .filter(key => to forall toPredicate(key))
+        .filter(key => from forall (key >=))
+        .filter(key => to forall (key <))
         .toList)
       filteredTree == keysIterator(newTree).toList
     }