Don't aggressively null out captured vars

Once they escape, we leave the references in the state
machines fields untouched.

2 files changed, 151 insertions, 7 deletions

diff --git a/src/test/scala/scala/async/TreeInterrogation.scala b/src/test/scala/scala/async/TreeInterrogation.scala
index 524e1a2..c8fe2d6 100644
--- a/src/test/scala/scala/async/TreeInterrogation.scala
+++ b/src/test/scala/scala/async/TreeInterrogation.scala
@@ -66,13 +66,43 @@ object TreeInterrogation extends App {
   withDebug {
     val cm = reflect.runtime.currentMirror
     val tb = mkToolbox("-cp ${toolboxClasspath} -Xprint:typer -uniqid")
-    import scala.async.Async._
+    import scala.async.internal.AsyncTestLV._
     val tree = tb.parse(
-      """ import _root_.scala.async.internal.AsyncId.{async, await}
+      """
+        | import scala.async.internal.AsyncTestLV._
+        | import scala.async.internal.AsyncTestLV
+        |
+        | case class MCell[T](var v: T)
+        | val f = async { MCell(1) }
+        |
+        | def m1(x: MCell[Int], y: Int): Int =
+        |   async { x.v + y }
+        | case class Cell[T](v: T)
+        |
         | async {
-        |   implicit def view(a: Int): String = ""
-        |   await(0).length
+        |   // state #1
+        |   val a: MCell[Int] = await(f)     // await$13$1
+        |   // state #2
+        |   var y = MCell(0)
+        |
+        |   while (a.v < 10) {
+        |     // state #4
+        |     a.v = a.v + 1
+        |     y = MCell(await(a).v + 1)      // await$14$1
+        |     // state #7
+        |   }
+        |
+        |   // state #3
+        |   assert(AsyncTestLV.log.exists(entry => entry._1 == "await$14$1"))
+        |
+        |   val b = await(m1(a, y.v))        // await$15$1
+        |   // state #8
+        |   assert(AsyncTestLV.log.exists(_ == ("a$1" -> MCell(10))))
+        |   assert(AsyncTestLV.log.exists(_ == ("y$1" -> MCell(11))))
+        |   b
         | }
+        |
+        |
         | """.stripMargin)
     println(tree)
     val tree1 = tb.typeCheck(tree.duplicate)
diff --git a/src/test/scala/scala/async/run/live/LiveVariablesSpec.scala b/src/test/scala/scala/async/run/live/LiveVariablesSpec.scala
index be62ed8..7d62f80 100644
--- a/src/test/scala/scala/async/run/live/LiveVariablesSpec.scala
+++ b/src/test/scala/scala/async/run/live/LiveVariablesSpec.scala
@@ -19,6 +19,7 @@ case class MCell[T](var v: T)
 
 
 class LiveVariablesSpec {
+  AsyncTestLV.clear()
 
   @Test
   def `zero out fields of reference type`() {
@@ -35,7 +36,7 @@ class LiveVariablesSpec {
       // a == Cell(1)
       val b: Cell[Int] = await(m1(a))  // await$2$1
       // b == Cell(2)
-      assert(AsyncTestLV.log.exists(_ == ("await$1$1" -> Cell(1))))
+      assert(AsyncTestLV.log.exists(_ == ("await$1$1" -> Cell(1))), AsyncTestLV.log)
       val res = await(m2(b))           // await$3$1
       assert(AsyncTestLV.log.exists(_ == ("await$2$1" -> Cell(2))))
       res
@@ -141,12 +142,125 @@ class LiveVariablesSpec {
 
       val b = await(m1(a, y.v))        // await$15$1
       // state #8
-      assert(AsyncTestLV.log.exists(_ == ("a$1" -> MCell(10))))
+      assert(AsyncTestLV.log.exists(_ == ("a$1" -> MCell(10))), AsyncTestLV.log)
       assert(AsyncTestLV.log.exists(_ == ("y$1" -> MCell(11))))
       b
     }
 
-    assert(m3() == 21)
+    assert(m3() == 21, m3())
   }
 
+  @Test
+  def `don't zero captured fields captured lambda`() {
+    val f = async {
+      val x = "x"
+      val y = "y"
+      await(0)
+      y.reverse
+      val f = () => assert(x != null)
+      await(0)
+      f
+    }
+    AsyncTestLV.assertNotNulledOut("x")
+    AsyncTestLV.assertNulledOut("y")
+    f()
+  }
+
+  @Test
+  def `don't zero captured fields captured by-name`() {
+    def func0[A](a: => A): () => A =  () => a
+    val f = async {
+      val x = "x"
+      val y = "y"
+      await(0)
+      y.reverse
+      val f = func0(assert(x != null))
+      await(0)
+      f
+    }
+    AsyncTestLV.assertNotNulledOut("x")
+    AsyncTestLV.assertNulledOut("y")
+    f()
+  }
+
+  @Test
+  def `don't zero captured fields nested class`() {
+    def func0[A](a: => A): () => A = () => a
+    val f = async {
+      val x = "x"
+      val y = "y"
+      await(0)
+      y.reverse
+      val f = new Function0[Unit] {
+        def apply = assert(x != null)
+      }
+      await(0)
+      f
+    }
+    AsyncTestLV.assertNotNulledOut("x")
+    AsyncTestLV.assertNulledOut("y")
+    f()
+  }
+
+  @Test
+  def `don't zero captured fields nested object`() {
+    def func0[A](a: => A): () => A = () => a
+    val f = async {
+      val x = "x"
+      val y = "y"
+      await(0)
+      y.reverse
+      object f extends Function0[Unit] {
+        def apply = assert(x != null)
+      }
+      await(0)
+      f
+    }
+    AsyncTestLV.assertNotNulledOut("x")
+    AsyncTestLV.assertNulledOut("y")
+    f()
+  }
+
+  @Test
+  def `don't zero captured fields nested def`() {
+    val f = async {
+      val x = "x"
+      val y = "y"
+      await(0)
+      y.reverse
+      def xx = x
+      val f = xx _
+      await(0)
+      f
+    }
+    AsyncTestLV.assertNotNulledOut("x")
+    AsyncTestLV.assertNulledOut("y")
+    f()
+  }
+
+  @Test
+  def `capture bug`() {
+    sealed trait Base
+    case class B1() extends Base
+    case class B2() extends Base
+    val outer = List[(Base, Int)]((B1(), 8))
+
+    def getMore(b: Base) = 4
+
+    def baz = async {
+      outer.head match {
+        case (a @ B1(), r) => {
+          val ents = await(getMore(a))
+
+          { () =>
+            println(a)
+            assert(a ne null)
+          }
+        }
+        case (b @ B2(), x) =>
+          () => ???
+      }
+    }
+    baz()
+  }
 }