Synthesize an object implementing Function#apply and resume()

There is a fly in the ointment: any This() trees in the provided
code are getting rebound to this class. Sounds like more
resetAttrs fun ahead.

The object also serves as the Future {} thunk.

5 files changed, 84 insertions, 52 deletions

diff --git a/src/main/scala/scala/async/Async.scala b/src/main/scala/scala/async/Async.scala
index f868f79..049dba0 100644
--- a/src/main/scala/scala/async/Async.scala
+++ b/src/main/scala/scala/async/Async.scala
@@ -106,39 +106,55 @@ abstract class AsyncBase {
         DefDef(mods, renameMap(dd.symbol), tparams, vparamss, tpt, c.resetAllAttrs(utils.substituteNames(rhs, renameMap)))
     }
 
-    val onCompleteHandler = asyncBlock.onCompleteHandler
+    val onCompleteHandler = {
+      Function(
+        List(ValDef(Modifiers(Flag.PARAM), name.tr, TypeTree(defn.TryAnyType), EmptyTree)),
+        asyncBlock.onCompleteHandler)
+    }
     val resumeFunTree = asyncBlock.resumeFunTree[T]
 
-    val prom: Expr[futureSystem.Prom[T]] = reify {
-      // Create the empty promise
-      val result$async = futureSystemOps.createProm[T].splice
-      // Initialize the state
-      var state$async = 0
-      // Resolve the execution context
-      val execContext$async = futureSystemOps.execContext.splice
-      var onCompleteHandler$async: util.Try[Any] => Unit = null
-
-      // Spawn a future to:
-      futureSystemOps.future[Unit] {
-        c.Expr[Unit](Block(
-          // define vars for all intermediate results that are accessed from multiple states
-          localVarTrees :+
-            // define the resume() method
-            resumeFunTree :+
-            // assign onComplete function. (The var breaks the circular dependency with resume)`
-            Assign(Ident(name.onCompleteHandler), onCompleteHandler),
-          // and get things started by calling resume()
-          Apply(Ident(name.resume), Nil)))
-      }(c.Expr[futureSystem.ExecContext](Ident(name.execContext))).splice
-      // Return the promise from this reify block...
-      result$async
+    val stateMachine: ModuleDef = {
+      val body: List[Tree] = {
+        val constr = DefDef(NoMods, nme.CONSTRUCTOR, List(), List(List()), TypeTree(), Block(List(Apply(Select(Super(This(tpnme.EMPTY), tpnme.EMPTY), nme.CONSTRUCTOR), List())), c.literalUnit.tree))
+        val stateVar = ValDef(Modifiers(Flag.MUTABLE), name.state, TypeTree(definitions.IntTpe), Literal(Constant(0)))
+        val result = ValDef(NoMods, name.result, TypeTree(), futureSystemOps.createProm[T].tree)
+        val execContext = ValDef(NoMods, name.execContext, TypeTree(), futureSystemOps.execContext.tree)
+        val applyDefDef: DefDef = {
+          val applyVParamss = List(List(ValDef(Modifiers(Flag.PARAM), name.tr, TypeTree(defn.TryAnyType), EmptyTree)))
+          val applyBody = asyncBlock.onCompleteHandler
+          DefDef(NoMods, name.apply, Nil, applyVParamss, TypeTree(definitions.UnitTpe), applyBody)
+        }
+        val apply0DefDef: DefDef = {
+          // We extend () => Unit so we can pass this class as the by-name argument to `Future.apply`.
+          // See SI-1247 for the the optimization that avoids creation of another thunk class.
+          val applyVParamss = List(List(ValDef(Modifiers(Flag.PARAM), name.tr, TypeTree(defn.TryAnyType), EmptyTree)))
+          val applyBody = asyncBlock.onCompleteHandler
+          DefDef(NoMods, name.apply, Nil, Nil, TypeTree(definitions.UnitTpe), Apply(Ident(name.resume), Nil))
+        }
+        List(constr, stateVar, result, execContext) ++ localVarTrees ++ List(resumeFunTree, applyDefDef, apply0DefDef)
+      }
+      val template = {
+        val `Try[Any] => Unit` = AppliedTypeTree(Ident(c.mirror.staticClass("scala.runtime.AbstractFunction1")), List(TypeTree(defn.TryAnyType), TypeTree(definitions.UnitTpe)))
+        val `() => Unit` = AppliedTypeTree(Ident(c.mirror.staticClass("scala.Function0")), List(TypeTree(definitions.UnitTpe)))
+        Template(List(`Try[Any] => Unit`, `() => Unit`), emptyValDef, body)
+      }
+      ModuleDef(NoMods, name.stateMachine, template)
     }
-    // ... and return its Future from the macro.
-    val result = futureSystemOps.promiseToFuture(prom)
 
-    AsyncUtils.vprintln(s"async state machine transform expands to:\n ${result.tree}")
+    def selectStateMachine(selection: TermName) = Select(Ident(name.stateMachine), selection)
+
+    val code = c.Expr[futureSystem.Fut[T]](Block(List[Tree](
+      stateMachine,
+      futureSystemOps.future(
+        c.Expr[Unit](Apply(selectStateMachine(name.apply), Nil)))
+        (c.Expr[futureSystem.ExecContext](selectStateMachine(name.execContext))).tree),
+      futureSystemOps.promiseToFuture(
+        c.Expr[futureSystem.Prom[T]](selectStateMachine(name.result))).tree
+    ))
+
+    AsyncUtils.vprintln(s"async state machine transform expands to:\n ${code.tree}")
+    code
 
-    result
   }
 
   def logDiagnostics(c: Context)(anfTree: c.Tree, states: Seq[String]) {
diff --git a/src/main/scala/scala/async/ExprBuilder.scala b/src/main/scala/scala/async/ExprBuilder.scala
index f8065f2..0655314 100644
--- a/src/main/scala/scala/async/ExprBuilder.scala
+++ b/src/main/scala/scala/async/ExprBuilder.scala
@@ -70,7 +70,7 @@ private[async] final case class ExprBuilder[C <: Context, FS <: FutureSystem](c:
 
     override def mkHandlerCaseForState: CaseDef = {
       val callOnComplete = futureSystemOps.onComplete(c.Expr(awaitable.expr),
-        c.Expr(Ident(name.onCompleteHandler)), c.Expr(Ident(name.execContext))).tree
+        c.Expr(This(tpnme.EMPTY)), c.Expr(Ident(name.execContext))).tree
       mkHandlerCase(state, stats :+ callOnComplete)
     }
 
@@ -310,20 +310,16 @@ private[async] final case class ExprBuilder[C <: Context, FS <: FutureSystem](c:
       val initStates = asyncStates.init
 
       /**
-       * lazy val onCompleteHandler = (tr: Try[Any]) => state match {
+       * // assumes tr: Try[Any] is in scope.
+       * //
+       * state match {
        * case 0 => {
        * x11 = tr.get.asInstanceOf[Double];
        * state = 1;
        * resume()
        * }
        */
-      val onCompleteHandler: Tree = {
-        val onCompleteHandlers = initStates.flatMap(_.mkOnCompleteHandler).toList
-        Function(
-          List(ValDef(Modifiers(Flag.PARAM), name.tr, TypeTree(defn.TryAnyType), EmptyTree)),
-          Match(Ident(name.state), onCompleteHandlers))
-      }
-
+      val onCompleteHandler: Tree = Match(Ident(name.state), initStates.flatMap(_.mkOnCompleteHandler).toList)
       /**
        * def resume(): Unit = {
        * try {
diff --git a/src/main/scala/scala/async/TransformUtils.scala b/src/main/scala/scala/async/TransformUtils.scala
index c684ea7..c66f874 100644
--- a/src/main/scala/scala/async/TransformUtils.scala
+++ b/src/main/scala/scala/async/TransformUtils.scala
@@ -18,19 +18,17 @@ private[async] final case class TransformUtils[C <: Context](c: C) {
 
     def suffixedName(prefix: String) = newTermName(suffix(prefix))
 
-    val state       = suffixedName("state")
-    val result      = suffixedName("result")
-    val resume      = suffixedName("resume")
-    val execContext = suffixedName("execContext")
-
-    // TODO do we need to freshen any of these?
-    val tr                = newTermName("tr")
-    val onCompleteHandler = suffixedName("onCompleteHandler")
-
-    val matchRes   = "matchres"
-    val ifRes      = "ifres"
-    val await      = "await"
-    val bindSuffix = "$bind"
+    val state        = suffixedName("state")
+    val result       = suffixedName("result")
+    val resume       = suffixedName("resume")
+    val execContext  = suffixedName("execContext")
+    val stateMachine = suffixedName("stateMachine")
+    val apply        = newTermName("apply")
+    val tr           = newTermName("tr")
+    val matchRes     = "matchres"
+    val ifRes        = "ifres"
+    val await        = "await"
+    val bindSuffix   = "$bind"
 
     def fresh(name: TermName): TermName = newTermName(fresh(name.toString))
 
diff --git a/src/test/scala/scala/async/TreeInterrogation.scala b/src/test/scala/scala/async/TreeInterrogation.scala
index ca4a309..4bdb84d 100644
--- a/src/test/scala/scala/async/TreeInterrogation.scala
+++ b/src/test/scala/scala/async/TreeInterrogation.scala
@@ -30,13 +30,14 @@ class TreeInterrogation {
     import tb.mirror.universe._
     val functions = tree1.collect {
       case f: Function => f
+      case t: Template => t
     }
     functions.size mustBe 1
 
     val varDefs = tree1.collect {
       case ValDef(mods, name, _, _) if mods.hasFlag(Flag.MUTABLE) => name
     }
-    varDefs.map(_.decoded).toSet mustBe (Set("state$async", "onCompleteHandler$async", "await$1", "await$2"))
+    varDefs.map(_.decoded.trim).toSet mustBe (Set("state$async", "await$1", "await$2"))
   }
 }
 
@@ -56,7 +57,15 @@ object TreeInterrogation extends App {
     val tb = mkToolbox("-cp target/scala-2.10/classes -Xprint:all")
     val tree = tb.parse(
       """ import _root_.scala.async.AsyncId._
-        | async { val a = 0; val x = await(a) - 1; def foo(z: Any) = (a.toDouble, x.toDouble, z); foo(await(2)) }
+        | object Test {
+        |   def blerg = 1
+        |   def check() {
+        |     async {
+        |       assert(this.blerg == 1)
+        |       assert(this == Test, this.getClass)
+        |     }
+        |   }
+        | }
         | """.stripMargin)
     println(tree)
     val tree1 = tb.typeCheck(tree.duplicate)
diff --git a/src/test/scala/scala/async/run/hygiene/Hygiene.scala b/src/test/scala/scala/async/run/hygiene/Hygiene.scala
index 0da2a4d..2aaf515 100644
--- a/src/test/scala/scala/async/run/hygiene/Hygiene.scala
+++ b/src/test/scala/scala/async/run/hygiene/Hygiene.scala
@@ -87,4 +87,17 @@ class HygieneSpec {
     }
     ext mustBe (14)
   }
+
+//  @Test def `this reference is maintained`() {
+//    object Test {
+//      def blerg = 1
+//      def check() {
+//        AsyncId.async {
+//          assert(this.blerg == 1)
+//          assert(this == Test, this.getClass)
+//        }
+//      }
+//    }
+//    Test.check()
+//  }
 }