From 00b5cb84df493aace270674054d2f6ddf3721131 Mon Sep 17 00:00:00 2001 From: Erik Rozendaal Date: Sun, 15 Jan 2012 13:48:00 +0100 Subject: Optimized implementation of TreeMap/TreeSet#to method. Performance of `to` and `until` is now the same. --- .../scala/collection/immutable/RedBlackTree.scala | 18 ++++++++----- .../scala/collection/immutable/TreeMap.scala | 10 ++++--- .../scala/collection/immutable/TreeSet.scala | 6 ++++- test/files/scalacheck/redblacktree.scala | 31 +++++++++++++--------- test/files/scalacheck/treemap.scala | 18 +++++++++++++ test/files/scalacheck/treeset.scala | 18 +++++++++++++ 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 -- cgit v1.2.3