Optimized implementation of TreeMap/TreeSet#to method.

Performance of `to` and `until` is now the same.

6 files changed, 77 insertions, 24 deletions

diff --git a/src/library/scala/collection/immutable/RedBlackTree.scala b/src/library/scala/collection/immutable/RedBlackTree.scala
index ebd88ce3fe..d8caeab096 100644
--- a/src/library/scala/collection/immutable/RedBlackTree.scala
+++ b/src/library/scala/collection/immutable/RedBlackTree.scala
@@ -45,7 +45,11 @@ 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], from: Option[A], until: Option[A])(implicit ordering: Ordering[A]): Tree[A, B] = blacken(rng(tree, from, until))
+  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 smallest[A, B](tree: Tree[A, B]): Tree[A, B] = {
     if (tree eq null) throw new NoSuchElementException("empty map")
@@ -198,13 +202,13 @@ object RedBlackTree {
     else append(tree.left, tree.right)
   }
 
-  private[this] def rng[A, B](tree: Tree[A, B], from: Option[A], until: Option[A])(implicit ordering: Ordering[A]): Tree[A, B] = {
+  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] = {
     if (tree eq null) return null
-    if (from == None && until == None) return tree
-    if (from != None && ordering.lt(tree.key, from.get)) return rng(tree.right, from, until);
-    if (until != None && ordering.lteq(until.get, tree.key)) return rng(tree.left, from, until);
-    val newLeft = rng(tree.left, from, None)
-    val newRight = rng(tree.right, None, until)
+    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 ((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 2bb8a566c6..3eba64dca3 100644
--- a/src/library/scala/collection/immutable/TreeMap.scala
+++ b/src/library/scala/collection/immutable/TreeMap.scala
@@ -62,9 +62,13 @@ 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,until)
-    new TreeMap[A,B](ntree)
+  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 firstKey = RB.smallest(tree).key
diff --git a/src/library/scala/collection/immutable/TreeSet.scala b/src/library/scala/collection/immutable/TreeSet.scala
index 8b95358d1c..5dd80e87a4 100644
--- a/src/library/scala/collection/immutable/TreeSet.scala
+++ b/src/library/scala/collection/immutable/TreeSet.scala
@@ -151,7 +151,11 @@ 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, until)
+    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 firstKey = head
diff --git a/test/files/scalacheck/redblacktree.scala b/test/files/scalacheck/redblacktree.scala
index 34fa8eae8d..14538c2352 100644
--- a/test/files/scalacheck/redblacktree.scala
+++ b/test/files/scalacheck/redblacktree.scala
@@ -174,36 +174,41 @@ package scala.collection.immutable.redblacktree {
   object TestRange extends RedBlackTreeTest with RedBlackTreeInvariants  {
     import RB._
 
-    override type ModifyParm = (Option[Int], Option[Int])
+    override type ModifyParm = (Option[Int], Boolean, Option[Int], Boolean)
     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, optionalTo)
+    } yield (optionalFrom, fromInclusive, optionalTo, toInclusive)
 
     override def modify(tree: Tree[String, Int], parm: ModifyParm): Tree[String, Int] = {
       val from = parm._1 flatMap (nodeAt(tree, _) map (_._1))
-      val to = parm._2 flatMap (nodeAt(tree, _) map (_._1))
-      range(tree, from, to)
+      val to = parm._3 flatMap (nodeAt(tree, _) map (_._1))
+      range(tree, from, parm._2, to, parm._4)
     }
 
     property("range boundaries respected") = forAll(genInput) { case (tree, parm, newTree) =>
       val from = parm._1 flatMap (nodeAt(tree, _) map (_._1))
-      val to = parm._2 flatMap (nodeAt(tree, _) map (_._1))
-      ("lower boundary" |: (from forall ( key => iterator(newTree).map(_._1) forall (key <=)))) &&
-      ("upper boundary" |: (to forall ( key => iterator(newTree).map(_._1) forall (key >))))
+      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))))
     }
 
     property("range returns all elements") = forAll(genInput) { case (tree, parm, newTree) =>
       val from = parm._1 flatMap (nodeAt(tree, _) map (_._1))
-      val to = parm._2 flatMap (nodeAt(tree, _) map (_._1))
-      val filteredTree = (iterator(tree)
-        .map(_._1)
-        .filter(key => from forall (key >=))
-        .filter(key => to forall (key <))
+      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 filteredTree = (keysIterator(tree)
+        .filter(key => from forall fromPredicate(key))
+        .filter(key => to forall toPredicate(key))
         .toList)
-      filteredTree == iterator(newTree).map(_._1).toList
+      filteredTree == keysIterator(newTree).toList
     }
   }
 }
diff --git a/test/files/scalacheck/treemap.scala b/test/files/scalacheck/treemap.scala
index 7d5f94d58b..ba6d117fd4 100644
--- a/test/files/scalacheck/treemap.scala
+++ b/test/files/scalacheck/treemap.scala
@@ -111,6 +111,24 @@ object Test extends Properties("TreeMap") {
     prefix.forall(_._1 < 0) && suffix.forall(_._1 >= 0) && subject == prefix ++ suffix
   }
 
+  property("from is inclusive") = forAll { (subject: TreeMap[Int, String]) => subject.nonEmpty ==> {
+    val n = choose(0, subject.size - 1).sample.get
+    val from = subject.drop(n).firstKey
+    subject.from(from).firstKey == from && subject.from(from).forall(_._1 >= from)
+  }}
+
+  property("to is inclusive") = forAll { (subject: TreeMap[Int, String]) => subject.nonEmpty ==> {
+    val n = choose(0, subject.size - 1).sample.get
+    val to = subject.drop(n).firstKey
+    subject.to(to).lastKey == to && subject.to(to).forall(_._1 <= to)
+  }}
+
+  property("until is exclusive") = forAll { (subject: TreeMap[Int, String]) => subject.size > 1 ==> {
+    val n = choose(1, subject.size - 1).sample.get
+    val until = subject.drop(n).firstKey
+    subject.until(until).lastKey == subject.take(n).lastKey && subject.until(until).forall(_._1 <= until)
+  }}
+
   property("remove single") = forAll { (subject: TreeMap[Int, String]) => subject.nonEmpty ==> {
     val key = oneOf(subject.keys.toSeq).sample.get
     val removed = subject - key
diff --git a/test/files/scalacheck/treeset.scala b/test/files/scalacheck/treeset.scala
index 7f99aec77e..e6d1b50860 100644
--- a/test/files/scalacheck/treeset.scala
+++ b/test/files/scalacheck/treeset.scala
@@ -107,6 +107,24 @@ object Test extends Properties("TreeSet") {
     prefix.forall(_ < 0) && suffix.forall(_ >= 0) && subject == prefix ++ suffix
   }
 
+  property("from is inclusive") = forAll { (subject: TreeSet[Int]) => subject.nonEmpty ==> {
+    val n = choose(0, subject.size - 1).sample.get
+    val from = subject.drop(n).firstKey
+    subject.from(from).firstKey == from && subject.from(from).forall(_ >= from)
+  }}
+
+  property("to is inclusive") = forAll { (subject: TreeSet[Int]) => subject.nonEmpty ==> {
+    val n = choose(0, subject.size - 1).sample.get
+    val to = subject.drop(n).firstKey
+    subject.to(to).lastKey == to && subject.to(to).forall(_ <= to)
+  }}
+
+  property("until is exclusive") = forAll { (subject: TreeSet[Int]) => subject.size > 1 ==> {
+    val n = choose(1, subject.size - 1).sample.get
+    val until = subject.drop(n).firstKey
+    subject.until(until).lastKey == subject.take(n).lastKey && subject.until(until).forall(_ <= until)
+  }}
+
   property("remove single") = forAll { (subject: TreeSet[Int]) => subject.nonEmpty ==> {
     val element = oneOf(subject.toSeq).sample.get
     val removed = subject - element