diff options
author | Erik Rozendaal <erik@deler.org> | 2011-12-17 21:18:48 +0100 |
---|---|---|
committer | Erik Rozendaal <erik@deler.org> | 2011-12-28 13:12:33 +0100 |
commit | 88ed93063419f6d09026e0ae466fe530f69af551 (patch) | |
tree | 70313a46b5f48fbf9026e5fd4152b42217b219ee | |
parent | 540ad0223ae26c0deae250c3ace2092904290a8b (diff) | |
download | scala-88ed93063419f6d09026e0ae466fe530f69af551.tar.gz scala-88ed93063419f6d09026e0ae466fe530f69af551.tar.bz2 scala-88ed93063419f6d09026e0ae466fe530f69af551.zip |
Use custom implementation for iterating over RedBlack trees. Raw
performance is much better than '++' based iterator.
-rw-r--r-- | src/library/scala/collection/immutable/RedBlack.scala | 36 |
1 files changed, 31 insertions, 5 deletions
diff --git a/src/library/scala/collection/immutable/RedBlack.scala b/src/library/scala/collection/immutable/RedBlack.scala index 9906c9896e..097df54af2 100644 --- a/src/library/scala/collection/immutable/RedBlack.scala +++ b/src/library/scala/collection/immutable/RedBlack.scala @@ -149,11 +149,9 @@ abstract class RedBlack[A] extends Serializable { def smallest: NonEmpty[B] = if (left.isEmpty) this else left.smallest - def toStream: Stream[(A,B)] = - left.toStream ++ Stream((key,value)) ++ right.toStream + def toStream: Stream[(A,B)] = iterator.toStream - def iterator: Iterator[(A, B)] = - left.iterator ++ Iterator.single(Pair(key, value)) ++ right.iterator + def iterator: Iterator[(A, B)] = new TreeIterator(this) def foreach[U](f: (A, B) => U) { left foreach f @@ -286,6 +284,34 @@ abstract class RedBlack[A] extends Serializable { override val right: Tree[B]) extends NonEmpty[B] { def isBlack = true } -} + private[this] class TreeIterator[B](tree: NonEmpty[B]) extends Iterator[(A, B)] { + import collection.mutable.Stack + + override def hasNext: Boolean = !next.isEmpty + override def next: (A, B) = next match { + case Empty => + throw new NoSuchElementException("next on empty iterator") + case tree: NonEmpty[B] => + val result = (tree.key, tree.value) + addLeftMostBranchToPath(tree.right) + next = if (path.isEmpty) Empty else path.pop() + result + } + + @annotation.tailrec + private[this] def addLeftMostBranchToPath(tree: Tree[B]) { + tree match { + case Empty => + case tree: NonEmpty[B] => + path.push(tree) + addLeftMostBranchToPath(tree.left) + } + } + + private[this] val path: Stack[NonEmpty[B]] = Stack.empty[NonEmpty[B]] + addLeftMostBranchToPath(tree) + private[this] var next: Tree[B] = path.pop() + } +} |