Enables testing the resetting of lifted local variables

- Adds a hook that lets a derived macro insert additional code
  when zero-ing out a lifted field.
- Adds a variant of the `AsyncId` macro that logs zeroed-out fields.
- Adds a test using this mechanism

6 files changed, 81 insertions, 8 deletions

diff --git a/src/main/scala/scala/async/continuations/AsyncBaseWithCPSFallback.scala b/src/main/scala/scala/async/continuations/AsyncBaseWithCPSFallback.scala
index 1a6ac87..20f5cce 100644
--- a/src/main/scala/scala/async/continuations/AsyncBaseWithCPSFallback.scala
+++ b/src/main/scala/scala/async/continuations/AsyncBaseWithCPSFallback.scala
@@ -92,7 +92,7 @@ trait AsyncBaseWithCPSFallback extends internal.AsyncBase {
                                           (execContext: c.Expr[futureSystem.ExecContext]): c.Expr[futureSystem.Fut[T]] = {
     AsyncUtils.vprintln("AsyncBaseWithCPSFallback.asyncImpl")
 
-    val asyncMacro = AsyncMacro(c, futureSystem)
+    val asyncMacro = AsyncMacro(c, this)
 
     if (!asyncMacro.reportUnsupportedAwaits(body.tree.asInstanceOf[asyncMacro.global.Tree], report = fallbackEnabled))
       super.asyncImpl[T](c)(body)(execContext)   // no unsupported awaits
diff --git a/src/main/scala/scala/async/internal/AsyncBase.scala b/src/main/scala/scala/async/internal/AsyncBase.scala
index 3690c2d..e44c27f 100644
--- a/src/main/scala/scala/async/internal/AsyncBase.scala
+++ b/src/main/scala/scala/async/internal/AsyncBase.scala
@@ -6,6 +6,7 @@ package scala.async.internal
 
 import scala.reflect.internal.annotations.compileTimeOnly
 import scala.reflect.macros.Context
+import scala.reflect.api.Universe
 
 /**
  * A base class for the `async` macro. Subclasses must provide:
@@ -45,7 +46,7 @@ abstract class AsyncBase {
                                  (execContext: c.Expr[futureSystem.ExecContext]): c.Expr[futureSystem.Fut[T]] = {
     import c.universe._
 
-    val asyncMacro = AsyncMacro(c, futureSystem)
+    val asyncMacro = AsyncMacro(c, self)
 
     val code = asyncMacro.asyncTransform[T](
       body.tree.asInstanceOf[asyncMacro.global.Tree],
@@ -59,4 +60,7 @@ abstract class AsyncBase {
     AsyncUtils.vprintln(s"async state machine transform expands to:\n ${code}")
     c.Expr[futureSystem.Fut[T]](code)
   }
+
+  protected[async] def nullOut(u: Universe)(name: u.Expr[String], v: u.Expr[Any]): u.Expr[Unit] =
+    u.reify { () }
 }
diff --git a/src/main/scala/scala/async/internal/AsyncId.scala b/src/main/scala/scala/async/internal/AsyncId.scala
index b9d82e2..7f7807f 100644
--- a/src/main/scala/scala/async/internal/AsyncId.scala
+++ b/src/main/scala/scala/async/internal/AsyncId.scala
@@ -6,6 +6,7 @@ package scala.async.internal
 
 import language.experimental.macros
 import scala.reflect.macros.Context
+import scala.reflect.api.Universe
 import scala.reflect.internal.SymbolTable
 
 object AsyncId extends AsyncBase {
@@ -17,6 +18,23 @@ object AsyncId extends AsyncBase {
   def asyncIdImpl[T: c.WeakTypeTag](c: Context)(body: c.Expr[T]): c.Expr[T] = asyncImpl[T](c)(body)(c.literalUnit)
 }
 
+object AsyncTestLV extends AsyncBase {
+  lazy val futureSystem = IdentityFutureSystem
+  type FS = IdentityFutureSystem.type
+
+  def async[T](body: T) = macro asyncIdImpl[T]
+
+  def asyncIdImpl[T: c.WeakTypeTag](c: Context)(body: c.Expr[T]): c.Expr[T] = asyncImpl[T](c)(body)(c.literalUnit)
+
+  var log: List[(String, Any)] = List()
+
+  def apply(name: String, v: Any): Unit =
+    log ::= (name -> v)
+
+  protected[async] override def nullOut(u: Universe)(name: u.Expr[String], v: u.Expr[Any]): u.Expr[Unit] =
+    u.reify { scala.async.internal.AsyncTestLV(name.splice, v.splice) }
+}
+
 /**
  * A trivial implementation of [[FutureSystem]] that performs computations
  * on the current thread. Useful for testing.
diff --git a/src/main/scala/scala/async/internal/AsyncMacro.scala b/src/main/scala/scala/async/internal/AsyncMacro.scala
index 1c97ca7..ee49923 100644
--- a/src/main/scala/scala/async/internal/AsyncMacro.scala
+++ b/src/main/scala/scala/async/internal/AsyncMacro.scala
@@ -4,15 +4,18 @@ import scala.tools.nsc.Global
 import scala.tools.nsc.transform.TypingTransformers
 
 object AsyncMacro {
-  def apply(c: reflect.macros.Context, futureSystem0: FutureSystem): AsyncMacro = {
+  def apply(c: reflect.macros.Context, base: AsyncBase): AsyncMacro = {
     import language.reflectiveCalls
     val powerContext = c.asInstanceOf[c.type { val universe: Global; val callsiteTyper: universe.analyzer.Typer }]
     new AsyncMacro {
-      val global: powerContext.universe.type = powerContext.universe
+      val global: powerContext.universe.type   = powerContext.universe
       val callSiteTyper: global.analyzer.Typer = powerContext.callsiteTyper
-      val futureSystem: futureSystem0.type = futureSystem0
-      val futureSystemOps: futureSystem.Ops {val universe: global.type} = futureSystem0.mkOps(global)
-      val macroApplication: global.Tree = c.macroApplication.asInstanceOf[global.Tree]
+      val macroApplication: global.Tree        = c.macroApplication.asInstanceOf[global.Tree]
+      // This member is required by `AsyncTransform`:
+      val asyncBase: AsyncBase                 = base
+      // These members are required by `ExprBuilder`:
+      val futureSystem: FutureSystem           = base.futureSystem
+      val futureSystemOps: futureSystem.Ops {val universe: global.type} = futureSystem.mkOps(global)
     }
   }
 }
diff --git a/src/main/scala/scala/async/internal/AsyncTransform.scala b/src/main/scala/scala/async/internal/AsyncTransform.scala
index 5c84f66..18caea4 100644
--- a/src/main/scala/scala/async/internal/AsyncTransform.scala
+++ b/src/main/scala/scala/async/internal/AsyncTransform.scala
@@ -5,6 +5,8 @@ trait AsyncTransform {
 
   import global._
 
+  val asyncBase: AsyncBase
+
   def asyncTransform[T](body: Tree, execContext: Tree, cpsFallbackEnabled: Boolean)
                        (resultType: WeakTypeTag[T]): Tree = {
 
@@ -68,7 +70,13 @@ trait AsyncTransform {
     for ((state, flds) <- assignsOf) {
       val assigns = flds.map { fld =>
         val fieldSym = fld.symbol
-        Assign(gen.mkAttributedStableRef(fieldSym.owner.thisType, fieldSym), gen.mkZero(fieldSym.info))
+        val zero     = gen.mkZero(fieldSym.info)
+        Block(
+          List(
+            asyncBase.nullOut(global)(Expr[String](Literal(Constant(fieldSym.name.toString))), Expr[Any](zero)).tree
+          ),
+          Assign(gen.mkAttributedStableRef(fieldSym.owner.thisType, fieldSym), zero)
+        )
       }
       val asyncState = asyncBlock.asyncStates.find(_.state == state).get
       asyncState.stats = assigns ++ asyncState.stats
diff --git a/src/test/scala/scala/async/run/live/LiveVariablesSpec.scala b/src/test/scala/scala/async/run/live/LiveVariablesSpec.scala
new file mode 100644
index 0000000..2cecffa
--- /dev/null
+++ b/src/test/scala/scala/async/run/live/LiveVariablesSpec.scala
@@ -0,0 +1,40 @@
+/*
+ * Copyright (C) 2012-2013 Typesafe Inc. <http://www.typesafe.com>
+ */
+
+package scala.async
+package run
+package live
+
+import org.junit.Test
+
+import internal.AsyncTestLV
+import AsyncTestLV._
+
+class LiveVariablesSpec {
+
+  @Test
+  def liveVars1() {
+    val f = async { 1 }
+
+    def m1(x: Int): Int =
+      async { x + 1 }
+
+    def m2(x: Int): String =
+      async { x.toString }
+
+    def m3() = async {
+      val a = await(f)          // await$1$1
+      // a == 1
+      val b = await(m1(a))      // await$2$1
+      // b == 2
+      assert(AsyncTestLV.log.exists(_ == ("await$1$1" -> 0)))
+      val res = await(m2(b))    // await$3$1
+      assert(AsyncTestLV.log.exists(_ == ("await$2$1" -> 0)))
+      res
+    }
+
+    assert(m3() == "2")
+  }
+
+}