Liveness analysis to avoid memory retention issues

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

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 { () }
 }

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.

src/main/scala/scala/async/internal/AsyncMacro.scala

@@ -4,23 +4,26 @@ 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)
     }
   }
 }
 
 private[async] trait AsyncMacro
   extends TypingTransformers
   with AnfTransform with TransformUtils with Lifter
-  with ExprBuilder with AsyncTransform with AsyncAnalysis {
+  with ExprBuilder with AsyncTransform with AsyncAnalysis with LiveVariables {
 
   val global: Global
   val callSiteTyper: global.analyzer.Typer

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 = {
 
@@ -29,6 +31,7 @@ trait AsyncTransform {
 
     val stateMachineType = applied("scala.async.StateMachine", List(futureSystemOps.promType[T](uncheckedBoundsResultTag), futureSystemOps.execContextType))
 
+    // Create `ClassDef` of state machine with empty method bodies for `resume` and `apply`.
     val stateMachine: ClassDef = {
       val body: List[Tree] = {
         val stateVar = ValDef(Modifiers(Flag.MUTABLE | Flag.PRIVATE | Flag.LOCAL), name.state, TypeTree(definitions.IntTpe), Literal(Constant(0)))
@@ -42,24 +45,45 @@ trait AsyncTransform {
         }
         List(emptyConstructor, stateVar, result, execContextValDef) ++ List(resumeFunTreeDummyBody, applyDefDefDummyBody, apply0DefDef)
       }
-      val template = {
-        Template(List(stateMachineType), emptyValDef, body)
-      }
+
+      val template = Template(List(stateMachineType), emptyValDef, body)
+
       val t = ClassDef(NoMods, name.stateMachineT, Nil, template)
       callSiteTyper.typedPos(macroPos)(Block(t :: Nil, Literal(Constant(()))))
       t
     }
 
+    val stateMachineClass = stateMachine.symbol
     val asyncBlock: AsyncBlock = {
-      val symLookup = new SymLookup(stateMachine.symbol, applyDefDefDummyBody.vparamss.head.head.symbol)
+      val symLookup = new SymLookup(stateMachineClass, applyDefDefDummyBody.vparamss.head.head.symbol)
       buildAsyncBlock(anfTree, symLookup)
     }
 
     logDiagnostics(anfTree, asyncBlock.asyncStates.map(_.toString))
 
+    val liftedFields: List[Tree] = liftables(asyncBlock.asyncStates)
+
+    // live variables analysis
+    // the result map indicates in which states a given field should be nulled out
+    val assignsOf = fieldsToNullOut(asyncBlock.asyncStates, liftedFields)
+
+    for ((state, flds) <- assignsOf) {
+      val assigns = flds.map { fld =>
+        val fieldSym = fld.symbol
+        Block(
+          List(
+            asyncBase.nullOut(global)(Expr[String](Literal(Constant(fieldSym.name.toString))), Expr[Any](Ident(fieldSym))).tree
+          ),
+          Assign(gen.mkAttributedStableRef(fieldSym.owner.thisType, fieldSym), gen.mkZero(fieldSym.info))
+        )
+      }
+      val asyncState = asyncBlock.asyncStates.find(_.state == state).get
+      asyncState.stats = assigns ++ asyncState.stats
+    }
+
     def startStateMachine: Tree = {
       val stateMachineSpliced: Tree = spliceMethodBodies(
-        liftables(asyncBlock.asyncStates),
+        liftedFields,
         stateMachine,
         atMacroPos(asyncBlock.onCompleteHandler[T]),
         atMacroPos(asyncBlock.resumeFunTree[T].rhs)
@@ -96,9 +120,16 @@ trait AsyncTransform {
     states foreach (s => AsyncUtils.vprintln(s))
   }
 
-  def spliceMethodBodies(liftables: List[Tree], tree: Tree, applyBody: Tree,
-                         resumeBody: Tree): Tree = {
-
+  /**
+   *  Build final `ClassDef` tree of state machine class.
+   *
+   *  @param  liftables  trees of definitions that are lifted to fields of the state machine class
+   *  @param  tree       `ClassDef` tree of the state machine class
+   *  @param  applyBody  tree of onComplete handler (`apply` method)
+   *  @param  resumeBody RHS of definition tree of `resume` method
+   *  @return            transformed `ClassDef` tree of the state machine class
+   */
+  def spliceMethodBodies(liftables: List[Tree], tree: ClassDef, applyBody: Tree, resumeBody: Tree): Tree = {
     val liftedSyms = liftables.map(_.symbol).toSet
     val stateMachineClass = tree.symbol
     liftedSyms.foreach {
@@ -112,7 +143,7 @@ trait AsyncTransform {
     // Replace the ValDefs in the splicee with Assigns to the corresponding lifted
     // fields. Similarly, replace references to them with references to the field.
     //
-    // This transform will be only be run on the RHS of `def foo`.
+    // This transform will only be run on the RHS of `def foo`.
     class UseFields extends MacroTypingTransformer {
       override def transform(tree: Tree): Tree = tree match {
         case _ if currentOwner == stateMachineClass          =>
@@ -150,6 +181,7 @@ trait AsyncTransform {
     }
     val treeSubst = tree
 
+    /* Fixes up DefDef: use lifted fields in `body` */
     def fixup(dd: DefDef, body: Tree, ctx: analyzer.Context): Tree = {
       val spliceeAnfFixedOwnerSyms = body
       val useField = new UseFields()
@@ -171,6 +203,7 @@ trait AsyncTransform {
         (ctx: analyzer.Context) =>
           val typedTree = fixup(dd, changeOwner(applyBody, callSiteTyper.context.owner, dd.symbol), ctx)
           typedTree
+
       case dd@DefDef(_, name.resume, _, _, _, _) if dd.symbol.owner == stateMachineClass =>
         (ctx: analyzer.Context) =>
           val changed = changeOwner(resumeBody, callSiteTyper.context.owner, dd.symbol)

src/main/scala/scala/async/internal/ExprBuilder.scala

@@ -25,11 +25,13 @@ trait ExprBuilder {
   trait AsyncState {
     def state: Int
 
+    def nextStates: List[Int]
+
     def mkHandlerCaseForState: CaseDef
 
     def mkOnCompleteHandler[T: WeakTypeTag]: Option[CaseDef] = None
 
-    def stats: List[Tree]
+    var stats: List[Tree]
 
     final def allStats: List[Tree] = this match {
       case a: AsyncStateWithAwait => stats :+ a.awaitable.resultValDef
@@ -43,9 +45,12 @@ trait ExprBuilder {
   }
 
   /** A sequence of statements that concludes with a unconditional transition to `nextState` */
-  final class SimpleAsyncState(val stats: List[Tree], val state: Int, nextState: Int, symLookup: SymLookup)
+  final class SimpleAsyncState(var stats: List[Tree], val state: Int, nextState: Int, symLookup: SymLookup)
     extends AsyncState {
 
+    def nextStates: List[Int] =
+      List(nextState)
+
     def mkHandlerCaseForState: CaseDef =
       mkHandlerCase(state, stats :+ mkStateTree(nextState, symLookup) :+ mkResumeApply(symLookup))
 
@@ -56,21 +61,24 @@ trait ExprBuilder {
   /** A sequence of statements with a conditional transition to the next state, which will represent
     * a branch of an `if` or a `match`.
     */
-  final class AsyncStateWithoutAwait(val stats: List[Tree], val state: Int) extends AsyncState {
+  final class AsyncStateWithoutAwait(var stats: List[Tree], val state: Int, val nextStates: List[Int]) extends AsyncState {
     override def mkHandlerCaseForState: CaseDef =
       mkHandlerCase(state, stats)
 
     override val toString: String =
-      s"AsyncStateWithoutAwait #$state"
+      s"AsyncStateWithoutAwait #$state, nextStates = $nextStates"
   }
 
   /** A sequence of statements that concludes with an `await` call. The `onComplete`
     * handler will unconditionally transition to `nextState`.
     */
-  final class AsyncStateWithAwait(val stats: List[Tree], val state: Int, nextState: Int,
+  final class AsyncStateWithAwait(var stats: List[Tree], val state: Int, nextState: Int,
                                   val awaitable: Awaitable, symLookup: SymLookup)
     extends AsyncState {
 
+    def nextStates: List[Int] =
+      List(nextState)
+
     override def mkHandlerCaseForState: CaseDef = {
       val callOnComplete = futureSystemOps.onComplete(Expr(awaitable.expr),
         Expr(This(tpnme.EMPTY)), Expr(Ident(name.execContext))).tree
@@ -147,7 +155,7 @@ trait ExprBuilder {
     def resultWithIf(condTree: Tree, thenState: Int, elseState: Int): AsyncState = {
       def mkBranch(state: Int) = Block(mkStateTree(state, symLookup) :: Nil, mkResumeApply(symLookup))
       this += If(condTree, mkBranch(thenState), mkBranch(elseState))
-      new AsyncStateWithoutAwait(stats.toList, state)
+      new AsyncStateWithoutAwait(stats.toList, state, List(thenState, elseState))
     }
 
     /**
@@ -169,12 +177,12 @@ trait ExprBuilder {
       }
       // 2. insert changed match tree at the end of the current state
       this += Match(scrutTree, newCases)
-      new AsyncStateWithoutAwait(stats.toList, state)
+      new AsyncStateWithoutAwait(stats.toList, state, caseStates)
     }
 
     def resultWithLabel(startLabelState: Int, symLookup: SymLookup): AsyncState = {
       this += Block(mkStateTree(startLabelState, symLookup) :: Nil, mkResumeApply(symLookup))
-      new AsyncStateWithoutAwait(stats.toList, state)
+      new AsyncStateWithoutAwait(stats.toList, state, List(startLabelState))
     }
 
     override def toString: String = {