Collapse all the onComplete handlers into a single function.

We're now down to two inner classes per async block:
the `onComplete` function, and the by-name argument
to the initial call to `future`.

Also caches the execution context in a val.

3 files changed, 68 insertions, 46 deletions

diff --git a/src/main/scala/scala/async/Async.scala b/src/main/scala/scala/async/Async.scala
index ac965b4..261279e 100644
--- a/src/main/scala/scala/async/Async.scala
+++ b/src/main/scala/scala/async/Async.scala
@@ -80,28 +80,51 @@ abstract class AsyncBase {
       case Block(stats, expr) =>
         val asyncBlockBuilder = new builder.AsyncBlockBuilder(stats, expr, 0, 1000, 1000, Map())
 
-        asyncBlockBuilder.asyncStates foreach vprintln
+        asyncBlockBuilder.asyncStates foreach (s => vprintln(s))
 
         val handlerCases: List[CaseDef] = asyncBlockBuilder.mkCombinedHandlerCases[T]()
 
-        val localVarTrees = asyncBlockBuilder.asyncStates.init.flatMap(_.allVarDefs).toList
+        val initStates = asyncBlockBuilder.asyncStates.init
+        val localVarTrees = initStates.flatMap(_.allVarDefs).toList
+
+        /*
+          lazy val onCompleteHandler = (tr: Try[Any]) => state match {
+            case 0 => {
+              x11 = tr.get.asInstanceOf[Double];
+              state = 1;
+              resume()
+            }
+            ...
+         */
+        val onCompleteHandler = {
+          val onCompleteHandlers = initStates.flatMap(_.mkOnCompleteHandler).toList
+          ValDef(Modifiers(LAZY), name.onCompleteHandler, TypeTree(),
+            Function(
+              List(ValDef(Modifiers(PARAM), name.tr, TypeTree(TryAnyType), EmptyTree)),
+              Match(Ident(name.state), onCompleteHandlers)))
+        }
 
         /*
           def resume(): Unit = {
             try {
-              state match { case 0 => ... ; ... }
+              state match {
+                case 0 => {
+                   f11 = exprReturningFuture
+                   f11.onComplete(onCompleteHandler)(context)
+                 }
+                ...
+               }
             } catch {
               case NonFatal(t) => result.failure(t)
             }
           }
          */
-        val nonFatalModule = builder.defn.NonFatalClass
         val resumeFunTree: c.Tree = DefDef(Modifiers(), name.resume, Nil, List(Nil), Ident(definitions.UnitClass),
           Try(
             Match(Ident(name.state), handlerCases),
             List(
               CaseDef(
-                Apply(Ident(nonFatalModule), List(Bind(name.tr, Ident(nme.WILDCARD)))),
+                Apply(Ident(NonFatalClass), List(Bind(name.tr, Ident(nme.WILDCARD)))),
                 EmptyTree,
                 Block(List({
                   val t = c.Expr[Throwable](Ident(name.tr))
@@ -110,17 +133,33 @@ abstract class AsyncBase {
 
 
         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
+          var execContext$async = futureSystemOps.execContext.splice
+
+          // Spawn a future to:
           futureSystemOps.future[Unit] {
             c.Expr[Unit](Block(
-              localVarTrees :+ resumeFunTree,
+              // define vars for all intermediate results
+              localVarTrees :+
+              // define the resume() method
+              resumeFunTree :+
+              // define the onComplete function
+              onCompleteHandler,
+              // and get things started by calling resume()
               Apply(Ident(name.resume), Nil)))
-          }(futureSystemOps.execContext).splice
+          }(c.Expr[futureSystem.ExecContext](Ident(name.execContext))).splice
+          // Return the promise from this reify block...
           result$async
         }
+        // ... and return its Future from the macro.
         val result = futureSystemOps.promiseToFuture(prom)
+
         vprintln(s"${c.macroApplication} \nexpands to:\n ${result.tree}")
+
         result
 
       case tree =>
diff --git a/src/main/scala/scala/async/AsyncUtils.scala b/src/main/scala/scala/async/AsyncUtils.scala
index 3b590be..77c155f 100644
--- a/src/main/scala/scala/async/AsyncUtils.scala
+++ b/src/main/scala/scala/async/AsyncUtils.scala
@@ -10,6 +10,6 @@ object AsyncUtils {
 
   private val verbose = false
   
-  private[async] def vprintln(s: Any): Unit = if (verbose)
+  private[async] def vprintln(s: => Any): Unit = if (verbose)
     println("[async] "+s)
 }
diff --git a/src/main/scala/scala/async/ExprBuilder.scala b/src/main/scala/scala/async/ExprBuilder.scala
index 3f82754..22f3738 100644
--- a/src/main/scala/scala/async/ExprBuilder.scala
+++ b/src/main/scala/scala/async/ExprBuilder.scala
@@ -27,10 +27,12 @@ final class ExprBuilder[C <: Context, FS <: FutureSystem](val c: C, val futureSy
     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 x1 = newTermName("x$1")
     val tr = newTermName("tr")
+    val onCompleteHandler = suffixedName("onCompleteHandler")
   }
 
   private val execContext = futureSystemOps.execContext
@@ -74,6 +76,21 @@ final class ExprBuilder[C <: Context, FS <: FutureSystem](val c: C, val futureSy
     def mkHandlerCaseForState(): CaseDef =
       mkHandlerCase(state, stats :+ mkStateTree(nextState) :+ mkResumeApply)
 
+    def mkOnCompleteHandler(): Option[CaseDef] = {
+      this match {
+        case aw: AsyncStateWithAwait =>
+          val tryGetTree =
+            Assign(
+              Ident(aw.resultName),
+              TypeApply(Select(Select(Ident(name.tr), Try_get), newTermName("asInstanceOf")), List(TypeTree(aw.resultType)))
+            )
+          val updateState = mkStateTree(nextState) // or increment?
+          Some(mkHandlerCase(state, List(tryGetTree, updateState, mkResumeApply)))
+        case _ =>
+          None
+      }
+    }
+
     def varDefForResult: Option[c.Tree] =
       None
 
@@ -106,48 +123,13 @@ final class ExprBuilder[C <: Context, FS <: FutureSystem](val c: C, val futureSy
     override val toString: String =
       s"AsyncStateWithAwait #$state, next = $nextState"
 
-    /* Make an `onComplete` invocation which increments the state upon resuming:
-     * 
-     *     awaitable.onComplete {
-     *       case tr =>
-     *         resultName = tr.get
-     *         state += 1
-     *         resume()
-     *     }
-     */
-    def mkOnCompleteIncrStateTree: c.Tree =
-      mkOnCompleteTree(mkInt_+(c.Expr[Int](Ident(name.state)))(c.literal(1)).tree)
-
-    /* Make an `onComplete` invocation which sets the state to `nextState` upon resuming:
-     * 
-     *     awaitable.onComplete {
-     *       case tr =>
-     *         resultName = tr.get
-     *         state = `nextState`
-     *         resume()
-     *     }
-     */
-    def mkOnCompleteStateTree(nextState: Int): c.Tree =
-      mkOnCompleteTree(c.literal(nextState).tree)
-
-    private def mkOnCompleteTree(nextState: Tree): c.Tree = {
-      val tryGetTree =
-        Assign(
-          Ident(resultName),
-          Select(Ident(name.tr), Try_get)
-        )
-
-      val updateState = mkStateTree(nextState)
-
-      val handlerTree =
-        Function(List(ValDef(Modifiers(PARAM), name.tr, TypeTree(tryType), EmptyTree)), Block(tryGetTree, updateState, mkResumeApply))
-
-      futureSystemOps.onComplete(c.Expr(awaitable), c.Expr(handlerTree), execContext).tree
+    private def mkOnCompleteTree: c.Tree = {
+      futureSystemOps.onComplete(c.Expr(awaitable), c.Expr(Ident(name.onCompleteHandler)), c.Expr(Ident(name.execContext))).tree
     }
 
     override def mkHandlerCaseForState(): CaseDef = {
       assert(awaitable != null)
-      mkHandlerCase(state, stats :+ mkOnCompleteIncrStateTree)
+      mkHandlerCase(state, stats :+ mkOnCompleteTree)
     }
 
     override def varDefForResult: Option[c.Tree] =
@@ -441,6 +423,7 @@ final class ExprBuilder[C <: Context, FS <: FutureSystem](val c: C, val futureSy
     val Try_get = methodSym(reify((null: scala.util.Try[Any]).get))
 
     val TryClass = c.mirror.staticClass("scala.util.Try")
+    val TryAnyType = appliedType(TryClass.toType, List(definitions.AnyTpe))
     val NonFatalClass = c.mirror.staticModule("scala.util.control.NonFatal")
 
     val Async_await = {