diff options
author | Paul Phillips <paulp@improving.org> | 2013-03-26 23:50:50 -0700 |
---|---|---|
committer | Paul Phillips <paulp@improving.org> | 2013-03-26 23:51:25 -0700 |
commit | e3ddb2d7dff859c9fb81d34d1c9687f72321a713 (patch) | |
tree | 1d520b035e5306214fe39ef5333ab7b55caa7867 /src | |
parent | 59d4998cf4a19eb5d44118d4867c2370e9a935e5 (diff) | |
download | scala-e3ddb2d7dff859c9fb81d34d1c9687f72321a713.tar.gz scala-e3ddb2d7dff859c9fb81d34d1c9687f72321a713.tar.bz2 scala-e3ddb2d7dff859c9fb81d34d1c9687f72321a713.zip |
Iterator.++ no longer blows the stack.
To my chagrin we still hadn't gotten this one. I took a new
approach which seems like a winner to me. Here's a benchmark:
object Test {
def run(n: Int) = println((1 to n).foldLeft(Iterator.empty: Iterator[Int])((res, _) => res ++ Iterator(1)) sum)
def main(args: Array[String]): Unit = run(args(0).toInt)
}
Runtime before this commit for various n:
500 0.403 real
1000 0.911 real
1500 2.351 real
2000 5.298 real
2500 10.184 real
Runtime after this commit, same n:
500 0.346 real
1000 0.359 real
1500 0.368 real
2000 0.379 real
2500 0.390 real
In the test case I dial it up to 100000.
Diffstat (limited to 'src')
-rw-r--r-- | src/library/scala/collection/Iterator.scala | 54 |
1 files changed, 36 insertions, 18 deletions
diff --git a/src/library/scala/collection/Iterator.scala b/src/library/scala/collection/Iterator.scala index c85a4fb6e7..72a23a0dd0 100644 --- a/src/library/scala/collection/Iterator.scala +++ b/src/library/scala/collection/Iterator.scala @@ -161,6 +161,41 @@ object Iterator { def hasNext = true def next = elem } + + /** Avoid stack overflows when applying ++ to lots of iterators by + * flattening the unevaluated iterators out into a vector of closures. + */ + private[scala] final class ConcatIterator[+A](initial: Vector[() => Iterator[A]]) extends Iterator[A] { + // current set to null when all iterators are exhausted + private[this] var current: Iterator[A] = Iterator.empty + private[this] var queue: Vector[() => Iterator[A]] = initial + // Advance current to the next non-empty iterator + private[this] def advance(): Boolean = { + if (queue.isEmpty) { + current = null + false + } + else { + current = queue.head() + queue = queue.tail + current.hasNext || advance() + } + } + def hasNext = (current ne null) && (current.hasNext || advance()) + def next() = if (hasNext) current.next else Iterator.empty.next + + override def ++[B >: A](that: => GenTraversableOnce[B]): Iterator[B] = + new ConcatIterator(queue :+ (() => that.toIterator)) + } + + private[scala] final class JoinIterator[+A](lhs: Iterator[A], that: => GenTraversableOnce[A]) extends Iterator[A] { + private[this] lazy val rhs: Iterator[A] = that.toIterator + def hasNext = lhs.hasNext || rhs.hasNext + def next = if (lhs.hasNext) lhs.next else rhs.next + + override def ++[B >: A](that: => GenTraversableOnce[B]) = + new ConcatIterator(Vector(() => this, () => that.toIterator)) + } } import Iterator.empty @@ -338,24 +373,7 @@ trait Iterator[+A] extends TraversableOnce[A] { * @usecase def ++(that: => Iterator[A]): Iterator[A] * @inheritdoc */ - def ++[B >: A](that: => GenTraversableOnce[B]): Iterator[B] = new AbstractIterator[B] { - // optimize a little bit to prevent n log n behavior. - private var cur : Iterator[B] = self - private var selfExhausted : Boolean = false - // since that is by-name, make sure it's only referenced once - - // if "val it = that" is inside the block, then hasNext on an empty - // iterator will continually reevaluate it. (ticket #3269) - lazy val it = that.toIterator - // the eq check is to avoid an infinite loop on "x ++ x" - def hasNext = cur.hasNext || (!selfExhausted && { - it.hasNext && { - cur = it - selfExhausted = true - true - } - }) - def next() = { hasNext; cur.next() } - } + def ++[B >: A](that: => GenTraversableOnce[B]): Iterator[B] = new Iterator.JoinIterator(self, that) /** Creates a new iterator by applying a function to all values produced by this iterator * and concatenating the results. |