Refactor some tree creation from Async to ExprBuilder.

3 files changed, 89 insertions, 72 deletions

diff --git a/src/main/scala/scala/async/Async.scala b/src/main/scala/scala/async/Async.scala
index 4c5c8f2..ef506a5 100644
--- a/src/main/scala/scala/async/Async.scala
+++ b/src/main/scala/scala/async/Async.scala
@@ -64,7 +64,6 @@ abstract class AsyncBase {
 
   def asyncImpl[T: c.WeakTypeTag](c: Context)(body: c.Expr[T]): c.Expr[futureSystem.Fut[T]] = {
     import c.universe._
-    import Flag._
 
     val builder = ExprBuilder[c.type, futureSystem.type](c, self.futureSystem)
     val anaylzer = AsyncAnalysis[c.type](c)
@@ -94,63 +93,17 @@ abstract class AsyncBase {
       }.toMap
     }
 
-    val startState = builder.stateAssigner.nextState()
-    val endState = Int.MaxValue
-
-    val asyncBlockBuilder = new builder.AsyncBlockBuilder(anfTree.stats, anfTree.expr, startState, endState, renameMap)
-    val handlerCases: List[CaseDef] = asyncBlockBuilder.mkCombinedHandlerCases[T]
-
-    import asyncBlockBuilder.asyncStates
+    val asyncBlock: builder.AsyncBlock = builder.build(anfTree, renameMap)
+    import asyncBlock.asyncStates
     logDiagnostics(c)(anfTree, asyncStates.map(_.toString))
-    val initStates = asyncStates.init
+
     val localVarTrees = anfTree.collect {
       case vd@ValDef(_, _, tpt, _) if renameMap contains vd.symbol =>
         utils.mkVarDefTree(tpt.tpe, renameMap(vd.symbol))
     }
 
-    /*
-      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
-      Function(
-        List(ValDef(Modifiers(PARAM), name.tr, TypeTree(defn.TryAnyType), EmptyTree)),
-        Match(Ident(name.state), onCompleteHandlers))
-    }
-
-    /*
-      def resume(): Unit = {
-        try {
-          state match {
-            case 0 => {
-               f11 = exprReturningFuture
-               f11.onComplete(onCompleteHandler)(context)
-             }
-            ...
-           }
-        } catch {
-          case NonFatal(t) => result.failure(t)
-        }
-      }
-     */
-    val resumeFunTree: c.Tree = DefDef(Modifiers(), name.resume, Nil, List(Nil), Ident(definitions.UnitClass),
-      Try(
-        Match(Ident(name.state), handlerCases),
-        List(
-          CaseDef(
-            Apply(Ident(defn.NonFatalClass), List(Bind(name.tr, Ident(nme.WILDCARD)))),
-            EmptyTree,
-            Block(List({
-              val t = c.Expr[Throwable](Ident(name.tr))
-              futureSystemOps.completeProm[T](c.Expr[futureSystem.Prom[T]](Ident(name.result)), reify(scala.util.Failure(t.splice))).tree
-            }), c.literalUnit.tree))), EmptyTree))
-
+    val onCompleteHandler = asyncBlock.onCompleteHandler
+    val resumeFunTree = asyncBlock.resumeFunTree[T]
 
     val prom: Expr[futureSystem.Prom[T]] = reify {
       // Create the empty promise
diff --git a/src/main/scala/scala/async/ExprBuilder.scala b/src/main/scala/scala/async/ExprBuilder.scala
index 00b06b1..5ae01f9 100644
--- a/src/main/scala/scala/async/ExprBuilder.scala
+++ b/src/main/scala/scala/async/ExprBuilder.scala
@@ -92,7 +92,7 @@ private[async] final case class ExprBuilder[C <: Context, FS <: FutureSystem](va
    */
   final class AsyncStateBuilder(state: Int, private val nameMap: Map[Symbol, c.Name]) {
     /* Statements preceding an await call. */
-    private val stats = ListBuffer[c.Tree]()
+    private val stats                      = ListBuffer[c.Tree]()
     /** The state of the target of a LabelDef application (while loop jump) */
     private var nextJumpState: Option[Int] = None
 
@@ -173,12 +173,12 @@ private[async] final case class ExprBuilder[C <: Context, FS <: FutureSystem](va
    * @param endState    the state to continue with
    * @param toRename    a `Map` for renaming the given key symbols to the mangled value names
    */
-  final class AsyncBlockBuilder(stats: List[c.Tree], expr: c.Tree, startState: Int, endState: Int,
-                                private val toRename: Map[Symbol, c.Name]) {
+  final private class AsyncBlockBuilder(stats: List[c.Tree], expr: c.Tree, startState: Int, endState: Int,
+                                        private val toRename: Map[Symbol, c.Name]) {
     val asyncStates = ListBuffer[AsyncState]()
 
-    private var stateBuilder = new AsyncStateBuilder(startState, toRename)
-    private var currState    = startState
+    var stateBuilder = new AsyncStateBuilder(startState, toRename)
+    var currState    = startState
 
     /* TODO Fall back to CPS plug-in if tree contains an `await` call. */
     def checkForUnsupportedAwait(tree: c.Tree) = if (tree exists {
@@ -186,7 +186,7 @@ private[async] final case class ExprBuilder[C <: Context, FS <: FutureSystem](va
       case _                             => false
     }) c.abort(tree.pos, "await must not be used in this position") //throw new FallbackToCpsException
 
-    private def nestedBlockBuilder(nestedTree: Tree, startState: Int, endState: Int) = {
+    def nestedBlockBuilder(nestedTree: Tree, startState: Int, endState: Int) = {
       val (nestedStats, nestedExpr) = statsAndExpr(nestedTree)
       new AsyncBlockBuilder(nestedStats, nestedExpr, startState, endState, toRename)
     }
@@ -262,22 +262,87 @@ private[async] final case class ExprBuilder[C <: Context, FS <: FutureSystem](va
     stateBuilder += expr
     val lastState = stateBuilder.resultSimple(endState)
     asyncStates += lastState
+  }
+
+  trait AsyncBlock {
+    def asyncStates: List[AsyncState]
+
+    def onCompleteHandler: Tree
+
+    def resumeFunTree[T]: Tree
+  }
+
+  def build(block: Block, toRename: Map[Symbol, c.Name]): AsyncBlock = {
+    val Block(stats, expr) = block
+    val startState = stateAssigner.nextState()
+    val endState = Int.MaxValue
 
-    def mkCombinedHandlerCases[T]: List[CaseDef] = {
-      val caseForLastState: CaseDef = {
-        val lastState = asyncStates.last
-        val lastStateBody = c.Expr[T](lastState.body)
-        val rhs = futureSystemOps.completeProm(
-          c.Expr[futureSystem.Prom[T]](Ident(name.result)), reify(scala.util.Success(lastStateBody.splice)))
-        mkHandlerCase(lastState.state, rhs.tree)
+    val blockBuilder = new AsyncBlockBuilder(stats, expr, startState, endState, toRename)
+
+    new AsyncBlock {
+      def asyncStates = blockBuilder.asyncStates.toList
+
+      def mkCombinedHandlerCases[T]: List[CaseDef] = {
+        val caseForLastState: CaseDef = {
+          val lastState = asyncStates.last
+          val lastStateBody = c.Expr[T](lastState.body)
+          val rhs = futureSystemOps.completeProm(
+            c.Expr[futureSystem.Prom[T]](Ident(name.result)), reify(scala.util.Success(lastStateBody.splice)))
+          mkHandlerCase(lastState.state, rhs.tree)
+        }
+        asyncStates.toList match {
+          case s :: Nil =>
+            List(caseForLastState)
+          case _        =>
+            val initCases = for (state <- asyncStates.toList.init) yield state.mkHandlerCaseForState
+            initCases :+ caseForLastState
+        }
       }
-      asyncStates.toList match {
-        case s :: Nil =>
-          List(caseForLastState)
-        case _        =>
-          val initCases = for (state <- asyncStates.toList.init) yield state.mkHandlerCaseForState
-          initCases :+ caseForLastState
+
+      val initStates = asyncStates.init
+
+      /**
+       * lazy val onCompleteHandler = (tr: Try[Any]) => 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))
       }
+
+      /**
+       * def resume(): Unit = {
+       * try {
+       * state match {
+       * case 0 => {
+       * f11 = exprReturningFuture
+       * f11.onComplete(onCompleteHandler)(context)
+       * }
+       * ...
+       * }
+       * } catch {
+       * case NonFatal(t) => result.failure(t)
+       * }
+       * }
+       */
+      def resumeFunTree[T]: Tree =
+        DefDef(Modifiers(), name.resume, Nil, List(Nil), Ident(definitions.UnitClass),
+          Try(
+            Match(Ident(name.state), mkCombinedHandlerCases[T]),
+            List(
+              CaseDef(
+                Apply(Ident(defn.NonFatalClass), List(Bind(name.tr, Ident(nme.WILDCARD)))),
+                EmptyTree,
+                Block(List({
+                  val t = c.Expr[Throwable](Ident(name.tr))
+                  futureSystemOps.completeProm[T](c.Expr[futureSystem.Prom[T]](Ident(name.result)), reify(scala.util.Failure(t.splice))).tree
+                }), c.literalUnit.tree))), EmptyTree))
     }
   }
 
diff --git a/src/main/scala/scala/async/TransformUtils.scala b/src/main/scala/scala/async/TransformUtils.scala
index 22099b5..8838bb3 100644
--- a/src/main/scala/scala/async/TransformUtils.scala
+++ b/src/main/scala/scala/async/TransformUtils.scala
@@ -24,7 +24,6 @@ private[async] final case class TransformUtils[C <: Context](val c: C) {
     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")