SI-8990 Allow GC during lazy evaluation of Stream#withFilter

  - The fact that StreamWithFilter was an inner class prevented
    collection of the head during processing after #withFilter,
    due to the reference to the outer class instance held in self.
  - By implementing StreamWithFilter outside of the Stream class,
    we gain control of the reference to the original Stream.
  - We clarify explicit "filter after first use" semantics for the
    reference to the original stream, which allows it to be GC'd
    during the processing of foreach, map, and flatMap.
  - Code is more DRY by implementing in terms of Stream#filter,
    which is already correctly lazy, and then Stream's #map,
    #flatMap, and #foreach, which already correctly allow GC.
  - Unfortunately, the type returned by Stream#withFilter *must*
    change, as it had previously inherited from the inner class
    TraversableLike#WithFilter, which is where the problematic
    reference to the outer class instance was held. Therefore,
    this change is targetted to 2.12.x. There doesn't appear to
    be any way to fix this without changing the type, sadly.
  - Special thanks to @paulp who suggested the likely cause of the
    issue at nescala 2015, and got me setup to build and run tests.
    Also thanks to @Ichoran and @retronym, who suggested that
    filter-after-first-use would preserve the reusable nature
    of the return value from #withFilter, rather than a single-shot
    null-after-first-use solution.

Review by @Ichoran and @retronym
Fixes #8990

1 files changed, 27 insertions, 52 deletions

diff --git a/src/library/scala/collection/immutable/Stream.scala b/src/library/scala/collection/immutable/Stream.scala
index b819302460..8be0b2fee2 100644
--- a/src/library/scala/collection/immutable/Stream.scala
+++ b/src/library/scala/collection/immutable/Stream.scala
@@ -524,57 +524,9 @@ self =>
    */
   override def filter(p: A => Boolean): Stream[A] = filterImpl(p, isFlipped = false) // This override is only left in 2.11 because of binary compatibility, see PR #3925
 
-  override final def withFilter(p: A => Boolean): StreamWithFilter = new StreamWithFilter(p)
-
-  /** A lazier implementation of WithFilter than TraversableLike's.
-   */
-  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()
-      }
-
-      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()
-      }
-
-      if (isStreamBuilder(bf)) asThat(tailFlatMap(Stream.this))
-      else super.flatMap(f)(bf)
-    }
-
-    override def foreach[B](f: A => B) =
-      for (x <- self)
-        if (p(x)) f(x)
-
-    override def withFilter(q: A => Boolean): StreamWithFilter =
-      new StreamWithFilter(x => p(x) && q(x))
-  }
+  /** A FilterMonadic which allows GC of the head of stream during processing */
+  @noinline // Workaround SI-9137, see https://github.com/scala/scala/pull/4284#issuecomment-73180791
+  override final def withFilter(p: A => Boolean): FilterMonadic[A, Stream[A]] = new Stream.StreamWithFilter(this, p)
 
   /** A lazier Iterator than LinearSeqLike's. */
   override def iterator: Iterator[A] = new StreamIterator(self)
@@ -1293,6 +1245,29 @@ object Stream extends SeqFactory[Stream] {
   private[immutable] def collectedTail[A, B, That](head: B, stream: Stream[A], pf: PartialFunction[A, B], bf: CanBuildFrom[Stream[A], B, That]) = {
     cons(head, stream.tail.collect(pf)(bf).asInstanceOf[Stream[B]])
   }
-}
 
+  /** An implementation of `FilterMonadic` allowing GC of the filtered-out elements of
+    * the `Stream` as it is processed.
+    *
+    * Because this is not an inner class of `Stream` with a reference to the original
+    * head, it is now possible for GC to collect any leading and filtered-out elements
+    * which do not satisfy the filter, while the tail is still processing (see SI-8990).
+    */
+  private[immutable] final class StreamWithFilter[A](sl: => Stream[A], p: A => Boolean) extends FilterMonadic[A, Stream[A]] {
+    private var s = sl                                              // set to null to allow GC after filtered
+    private lazy val filtered = { val f = s filter p; s = null; f } // don't set to null if throw during filter
+
+    def map[B, That](f: A => B)(implicit bf: CanBuildFrom[Stream[A], B, That]): That =
+      filtered map f
 
+    def flatMap[B, That](f: A => scala.collection.GenTraversableOnce[B])(implicit bf: CanBuildFrom[Stream[A], B, That]): That =
+      filtered flatMap f
+
+    def foreach[U](f: A => U): Unit =
+      filtered foreach f
+
+    def withFilter(q: A => Boolean): FilterMonadic[A, Stream[A]] =
+      new StreamWithFilter[A](filtered, q)
+  }
+
+}