Skip to content
Browse files

Make Stream.withFilter.{map,flatMap} run in constant stack space

The included test currently fails because `map` and `flatMap` do not
run in constant stack space on a stream returned by `Stream.withFilter`,
as I reported here:
https://groups.google.com/d/msg/scala-language/WqJR38REXnk/saaSiDdmyqoJ

Fix the problem and add a simple testcase.

Note that the stack space consumed when producing an element of this stream is
proportional to the number of elements failing the test before the next
success. The stack space consumed to produce the stream itself is the
space needed to produce the first element, that is, is proportional to
the number of failures before the first success.
  • Loading branch information...
1 parent 8e7f44c commit 53130d21434a67d6d9b443d531578a04eb9e510f @Blaisorblade Blaisorblade committed Aug 20, 2012
View
35 src/library/scala/collection/immutable/Stream.scala
@@ -479,22 +479,57 @@ self =>
final class StreamWithFilter(p: A => Boolean) extends WithFilter(p) {
override def map[B, That](f: A => B)(implicit bf: CanBuildFrom[Stream[A], B, That]): That = {
+ def tailMap(coll: Stream[A]): Stream[B] = {
+ var head: A = null.asInstanceOf[A]
+ var tail: Stream[A] = coll
+ while (true) {
+ if (tail.isEmpty)
+ return Stream.Empty
+ head = tail.head
+ tail = tail.tail
+ if (p(head))
+ return cons(f(head), tailMap(tail))
+ }
+ throw new RuntimeException()
+ }
+
+ /*
def tailMap = asStream[B](tail withFilter p map f)
if (isStreamBuilder(bf)) asThat(
if (isEmpty) Stream.Empty
else if (p(head)) cons(f(head), tailMap)
else tailMap
+ //XXX Alternative: what about having a Stream.step constructor?
)
+ */
+ if (isStreamBuilder(bf)) asThat(tailMap(Stream.this))
else super.map(f)(bf)
}
override def flatMap[B, That](f: A => GenTraversableOnce[B])(implicit bf: CanBuildFrom[Stream[A], B, That]): That = {
+ def tailFlatMap(coll: Stream[A]): Stream[B] = {
+ var head: A = null.asInstanceOf[A]
+ var tail: Stream[A] = coll
+ while (true) {
+ if (tail.isEmpty)
+ return Stream.Empty
+ head = tail.head
+ tail = tail.tail
+ if (p(head))
+ return f(head).toStream append tailFlatMap(tail)
+ }
+ throw new RuntimeException()
+ }
+
+ /*
def tailFlatMap = asStream[B](tail withFilter p flatMap f)
if (isStreamBuilder(bf)) asThat(
if (isEmpty) Stream.Empty
else if (p(head)) f(head).toStream append tailFlatMap
else tailFlatMap
)
+ */
+ if (isStreamBuilder(bf)) asThat(tailFlatMap(Stream.this))
else super.flatMap(f)(bf)
}
View
2 test/files/run/stream-stack-overflow-filter-map.check
@@ -0,0 +1,2 @@
+Stream()
+Stream()
View
14 test/files/run/stream-stack-overflow-filter-map.scala
@@ -0,0 +1,14 @@
+object Test extends App {
+ //This runs fine.
+ val resFMap1 = (1 to 10000).toStream filter (_ => false) flatMap (Seq(_))
+ val resMap1 = (1 to 10000).toStream filter (_ => false) map (_ + 1)
+ assert(resMap1.isEmpty)
+ assert(resFMap1.isEmpty)
+ println(resMap1)
+ println(resFMap1)
+ //This will cause a stack overflow
+ val resFMap2 = (1 to 10000).toStream withFilter (_ => false) flatMap (Seq(_))
+ val resMap2 = (1 to 10000).toStream withFilter (_ => false) map (_ + 1)
+ assert(resMap1 == resMap2)
+ assert(resFMap1 == resFMap2)
+}

0 comments on commit 53130d2

Please sign in to comment.
Something went wrong with that request. Please try again.