diff options
author | Erik Rozendaal <erik@deler.org> | 2012-01-04 17:10:20 +0100 |
---|---|---|
committer | Erik Rozendaal <erik@deler.org> | 2012-01-04 22:36:21 +0100 |
commit | 72ec0ac869a29fca9ea0d45a3f70f1e9e1babaaf (patch) | |
tree | 37591a9c2cea97da7afdf55ee03a84fa2133a1db /src/library | |
parent | 5c05f66b619ea9c8a543518fd9d7d601268c6f19 (diff) | |
download | scala-72ec0ac869a29fca9ea0d45a3f70f1e9e1babaaf.tar.gz scala-72ec0ac869a29fca9ea0d45a3f70f1e9e1babaaf.tar.bz2 scala-72ec0ac869a29fca9ea0d45a3f70f1e9e1babaaf.zip |
Optimize foreach and iterators.
Diffstat (limited to 'src/library')
-rw-r--r-- | src/library/scala/collection/immutable/RedBlack.scala | 108 | ||||
-rw-r--r-- | src/library/scala/collection/immutable/TreeMap.scala | 5 | ||||
-rw-r--r-- | src/library/scala/collection/immutable/TreeSet.scala | 2 |
3 files changed, 71 insertions, 44 deletions
diff --git a/src/library/scala/collection/immutable/RedBlack.scala b/src/library/scala/collection/immutable/RedBlack.scala index 2537d043fd..6af6b6ef03 100644 --- a/src/library/scala/collection/immutable/RedBlack.scala +++ b/src/library/scala/collection/immutable/RedBlack.scala @@ -11,6 +11,7 @@ package scala.collection package immutable +import annotation.tailrec import annotation.meta.getter /** An object containing the RedBlack tree implementation used by for `TreeMaps` and `TreeSets`. @@ -37,7 +38,7 @@ object RedBlack { case tree => Some(tree.value) } - @annotation.tailrec + @tailrec def lookup[A, B](tree: Tree[A, B], x: A)(implicit ordering: Ordering[A]): Tree[A, B] = if (tree eq null) null else { val cmp = ordering.compare(x, tree.key) if (cmp < 0) lookup(tree.left, x) @@ -64,18 +65,19 @@ object RedBlack { } def foreach[A, B, U](tree: Tree[A, B], f: ((A, B)) => U): Unit = if (tree ne null) { - foreach(tree.left, f) + if (tree.left ne null) foreach(tree.left, f) f((tree.key, tree.value)) - foreach(tree.right, f) + if (tree.right ne null) foreach(tree.right, f) } def foreachKey[A, U](tree: Tree[A, _], f: A => U): Unit = if (tree ne null) { - foreachKey(tree.left, f) + if (tree.left ne null) foreachKey(tree.left, f) f(tree.key) - foreachKey(tree.right, f) + if (tree.right ne null) foreachKey(tree.right, f) } - def iterator[A, B](tree: Tree[A, B]): Iterator[(A, B)] = if (tree eq null) Iterator.empty else new TreeIterator(tree) - def keyIterator[A, _](tree: Tree[A, _]): Iterator[A] = if (tree eq null) Iterator.empty else new TreeKeyIterator(tree) + def iterator[A, B](tree: Tree[A, B]): Iterator[(A, B)] = new EntriesIterator(tree) + def keysIterator[A, _](tree: Tree[A, _]): Iterator[A] = new KeysIterator(tree) + def valuesIterator[_, B](tree: Tree[_, B]): Iterator[B] = new ValuesIterator(tree) private[this] def balanceLeft[A, B, B1 >: B](isBlack: Boolean, z: A, zv: B, l: Tree[A, B1], d: Tree[A, B1]): Tree[A, B1] = { if (isRedTree(l) && isRedTree(l.left)) @@ -283,7 +285,7 @@ object RedBlack { @(inline @getter) final val left: Tree[A, B], @(inline @getter) final val right: Tree[A, B]) extends Serializable { - @(inline @getter) final val count: Int = 1 + RedBlack.count(left) + RedBlack.count(right) + final val count: Int = 1 + RedBlack.count(left) + RedBlack.count(right) def isBlack: Boolean def nth(n: Int): Tree[A, B] = { val count = RedBlack.count(left) @@ -322,53 +324,75 @@ object RedBlack { def unapply[A, B](t: BlackTree[A, B]) = Some((t.key, t.value, t.left, t.right)) } - private[this] class TreeIterator[A, B](tree: Tree[A, B]) extends Iterator[(A, B)] { + private[this] abstract class TreeIterator[A, B, R](tree: Tree[A, B]) extends Iterator[R] { + protected[this] def nextResult(tree: Tree[A, B]): R + override def hasNext: Boolean = next ne null - override def next: (A, B) = next match { + override def next: R = next match { case null => throw new NoSuchElementException("next on empty iterator") case tree => - addLeftMostBranchToPath(tree.right) - next = if (path.isEmpty) null else path.pop() - (tree.key, tree.value) + next = findNext(tree.right) + nextResult(tree) } - @annotation.tailrec - private[this] def addLeftMostBranchToPath(tree: Tree[A, B]) { - if (tree ne null) { - path.push(tree) - addLeftMostBranchToPath(tree.left) + @tailrec + private[this] def findNext(tree: Tree[A, B]): Tree[A, B] = { + if (tree eq null) popPath() + else if (tree.left eq null) tree + else { + pushPath(tree) + findNext(tree.left) } } - private[this] val path = mutable.ArrayStack.empty[Tree[A, B]] - addLeftMostBranchToPath(tree) - private[this] var next: Tree[A, B] = path.pop() - } - - private[this] class TreeKeyIterator[A](tree: Tree[A, _]) extends Iterator[A] { - override def hasNext: Boolean = next ne null - - override def next: A = next match { - case null => - throw new NoSuchElementException("next on empty iterator") - case tree => - addLeftMostBranchToPath(tree.right) - next = if (path.isEmpty) null else path.pop() - tree.key + private[this] def pushPath(tree: Tree[A, B]) { + try { + path(index) = tree + index += 1 + } catch { + case _: ArrayIndexOutOfBoundsException => + // Either the tree became unbalanced or we calculated the maximum height incorrectly. + // To avoid crashing the iterator we expand the path array. Obviously this should never + // happen... + // + // An exception handler is used instead of an if-condition to optimize the normal path. + assert(index >= path.length) + path :+= null + pushPath(tree) + } + } + private[this] def popPath(): Tree[A, B] = if (index == 0) null else { + index -= 1 + path(index) } - @annotation.tailrec - private[this] def addLeftMostBranchToPath(tree: Tree[A, _]) { - if (tree ne null) { - path.push(tree) - addLeftMostBranchToPath(tree.left) - } + private[this] var path = if (tree eq null) null else { + /* + * According to "Ralf Hinze. Constructing red-black trees" [http://www.cs.ox.ac.uk/ralf.hinze/publications/#P5] + * the maximum height of a red-black tree is 2*log_2(n + 2) - 2. + * + * According to {@see Integer#numberOfLeadingZeros} ceil(log_2(n)) = (32 - Integer.numberOfLeadingZeros(n - 1)) + * + * We also don't store the deepest nodes in the path so the maximum path length is further reduced by one. + */ + val maximumHeight = 2 * (32 - Integer.numberOfLeadingZeros(tree.count + 2 - 1)) - 2 - 1 + new Array[Tree[A, B]](maximumHeight) } + private[this] var index = 0 + private[this] var next: Tree[A, B] = findNext(tree) + } + + private[this] class EntriesIterator[A, B](tree: Tree[A, B]) extends TreeIterator[A, B, (A, B)](tree) { + override def nextResult(tree: Tree[A, B]) = (tree.key, tree.value) + } + + private[this] class KeysIterator[A, B](tree: Tree[A, B]) extends TreeIterator[A, B, A](tree) { + override def nextResult(tree: Tree[A, B]) = tree.key + } - private[this] val path = mutable.ArrayStack.empty[Tree[A, _]] - addLeftMostBranchToPath(tree) - private[this] var next: Tree[A, _] = path.pop() + private[this] class ValuesIterator[A, B](tree: Tree[A, B]) extends TreeIterator[A, B, B](tree) { + override def nextResult(tree: Tree[A, B]) = tree.value } } diff --git a/src/library/scala/collection/immutable/TreeMap.scala b/src/library/scala/collection/immutable/TreeMap.scala index 45e936444f..6e8cf625f4 100644 --- a/src/library/scala/collection/immutable/TreeMap.scala +++ b/src/library/scala/collection/immutable/TreeMap.scala @@ -196,7 +196,10 @@ class TreeMap[A, +B] private (tree: RedBlack.Tree[A, B])(implicit val ordering: * * @return the new iterator */ - def iterator: Iterator[(A, B)] = RB.iterator(tree) + override def iterator: Iterator[(A, B)] = RB.iterator(tree) + + override def keysIterator: Iterator[A] = RB.keysIterator(tree) + override def valuesIterator: Iterator[B] = RB.valuesIterator(tree) override def contains(key: A): Boolean = RB.contains(tree, key) override def isDefinedAt(key: A): Boolean = RB.contains(tree, key) diff --git a/src/library/scala/collection/immutable/TreeSet.scala b/src/library/scala/collection/immutable/TreeSet.scala index 00ebeab868..7c27e9f5b0 100644 --- a/src/library/scala/collection/immutable/TreeSet.scala +++ b/src/library/scala/collection/immutable/TreeSet.scala @@ -149,7 +149,7 @@ class TreeSet[A] private (tree: RedBlack.Tree[A, Unit])(implicit val ordering: O * * @return the new iterator */ - def iterator: Iterator[A] = RB.keyIterator(tree) + def iterator: Iterator[A] = RB.keysIterator(tree) override def foreach[U](f: A => U) = RB.foreachKey(tree, f) |