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.

3 files changed, 51 insertions, 0 deletions

diff --git a/src/library/scala/collection/immutable/Stream.scala b/src/library/scala/collection/immutable/Stream.scala
index 1bf1e20694..c8e7e7547f 100644
--- a/src/library/scala/collection/immutable/Stream.scala
+++ b/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)
     }
 
diff --git a/test/files/run/stream-stack-overflow-filter-map.check b/test/files/run/stream-stack-overflow-filter-map.check
new file mode 100644
index 0000000000..dbf856668b
--- /dev/null
+++ b/test/files/run/stream-stack-overflow-filter-map.check
@@ -0,0 +1,2 @@
+Stream()
+Stream()
diff --git a/test/files/run/stream-stack-overflow-filter-map.scala b/test/files/run/stream-stack-overflow-filter-map.scala
new file mode 100644
index 0000000000..e2bd47bfca
--- /dev/null
+++ b/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)
+}