Synthesize a single PartialFunction rather than two and an orElse call.

Give the handler a real class name.

2 files changed, 33 insertions, 30 deletions

diff --git a/src/main/scala/scala/async/Async.scala b/src/main/scala/scala/async/Async.scala
index 0a3bbbf..13cca71 100644
--- a/src/main/scala/scala/async/Async.scala
+++ b/src/main/scala/scala/async/Async.scala
@@ -23,7 +23,7 @@ object Async extends AsyncBase {
   def async[T](body: T) = macro asyncImpl[T]
 
   override def asyncImpl[T: c.WeakTypeTag](c: Context)(body: c.Expr[T]): c.Expr[Future[T]] = super.asyncImpl[T](c)(body)
-}
+    }
 
 object AsyncId extends AsyncBase {
   lazy val futureSystem = IdentityFutureSystem
@@ -82,29 +82,25 @@ abstract class AsyncBase extends AsyncUtils {
       case Block(stats, expr) =>
         val asyncBlockBuilder = new builder.AsyncBlockBuilder(stats, expr, 0, 1000, 1000, Map())
 
-        vprintln(s"states of current method:")
         asyncBlockBuilder.asyncStates foreach vprintln
 
-        val handlerExpr = asyncBlockBuilder.mkCombinedHandlerExpr()
-
-        vprintln(s"GENERATED handler expr:")
-        vprintln(handlerExpr)
+        val handlerCases: List[(CaseDef, Int)] = asyncBlockBuilder.mkCombinedHandlerCases()
 
-        val handlerForLastState: c.Expr[PartialFunction[Int, Unit]] = {
+        val caseForLastState: (CaseDef, Int) = {
           val lastState = asyncBlockBuilder.asyncStates.last
           val lastStateBody = c.Expr[T](lastState.body)
-          builder.mkHandler(lastState.state, futureSystemOps.completeProm(c.Expr[futureSystem.Prom[T]](Ident(name.result)), reify(scala.util.Success(lastStateBody.splice))))
+          val rhs = futureSystemOps.completeProm(c.Expr[futureSystem.Prom[T]](Ident(name.result)), reify(scala.util.Success(lastStateBody.splice)))
+          builder.mkHandlerCase(lastState.state, rhs.tree) -> lastState.state
         }
 
-        vprintln("GENERATED handler for last state:")
-        vprintln(handlerForLastState)
+        val combinedHander = c.Expr[PartialFunction[Int, Unit]](builder.mkHandlerTreeFor(handlerCases :+ caseForLastState))
 
         val localVarTrees = asyncBlockBuilder.asyncStates.init.flatMap(_.allVarDefs).toList
 
         /*
           def resume(): Unit = {
             try {
-              (handlerExpr.splice orElse handlerForLastState.splice)(state)
+              combinedHander(state)
             } catch {
               case NonFatal(t) => result.failure(t)
             }
@@ -114,8 +110,7 @@ abstract class AsyncBase extends AsyncUtils {
         val resumeFunTree: c.Tree = DefDef(Modifiers(), name.resume, List(), List(List()), Ident(definitions.UnitClass),
           Try(
             reify {
-              val combinedHandler = mkPartialFunction_orElse(handlerExpr)(handlerForLastState).splice
-              combinedHandler.apply(c.Expr[Int](Ident(name.state)).splice)
+              combinedHander.splice.apply(c.Expr[Int](Ident(name.state)).splice)
             }.tree
             ,
             List(
@@ -139,15 +134,9 @@ abstract class AsyncBase extends AsyncUtils {
           result$async
         }
         val result = futureSystemOps.promiseToFuture(prom)
-//        println(s"${c.macroApplication} \nexpands to:\n ${result.tree}")
-//        val positions = result.tree.collect {
-//          case t => (t.toString.take(10).replaceAll("\n", "\\n"), t.pos)
-//        }
-//        println(positions.mkString("\n"))
-
+        vprintln(s"${c.macroApplication} \nexpands to:\n ${result.tree}")
         result
 
-
       case tree =>
         c.abort(c.macroApplication.pos, s"expression not supported by async: ${tree}")
     }
diff --git a/src/main/scala/scala/async/ExprBuilder.scala b/src/main/scala/scala/async/ExprBuilder.scala
index b563869..6578df8 100644
--- a/src/main/scala/scala/async/ExprBuilder.scala
+++ b/src/main/scala/scala/async/ExprBuilder.scala
@@ -21,17 +21,19 @@ final class ExprBuilder[C <: Context, FS <: FutureSystem](val c: C, val futureSy
 
   object name {
     // TODO do we need to freshen any of these?
-    def expanded(prefix: String) = newTermName(prefix + "$async")
-    val resume = expanded("resume")
-    val state = expanded("state")
-    val result = expanded("result")
+    def suffix(string: String) = string + "$async"
+    def expandedTermName(prefix: String) = newTermName(suffix(prefix))
+    def expandedTypeName(prefix: String) = newTypeName(suffix(prefix))
+    val resume = expandedTermName("resume")
+    val state = expandedTermName("state")
+    val result = expandedTermName("result")
     val tr = newTermName("tr")
     val any = newTermName("any")
     val x1 = newTermName("x$1")
     val apply = newTermName("apply")
     val isDefinedAt = newTermName("isDefinedAt")
 
-    val anon = newTypeName("$anon")
+    val asyncHander = expandedTypeName("Handler")
   }
 
   private val execContext = futureSystemOps.execContext
@@ -98,31 +100,36 @@ final class ExprBuilder[C <: Context, FS <: FutureSystem](val c: C, val futureSy
       rhs
     )
 
+  private def paramValDef(name: TermName, sym: Symbol) = ValDef(Modifiers(PARAM), name, Ident(sym), EmptyTree)
+  private def paramValDef(name: TermName, tpe: Type) = ValDef(Modifiers(PARAM), name, TypeTree(tpe), EmptyTree)
+
   def mkHandlerTreeFor(cases: List[(CaseDef, Int)]): c.Tree = {
     val partFunIdent = Ident(defn.PartialFunctionClass)
+//    val partFunTpe = appliedType(defn.PartialFunctionClass.tpe, definitions.IntTpe, definitions.UnitTpe)
     val intIdent = Ident(definitions.IntClass)
     val unitIdent = Ident(definitions.UnitClass)
 
     val caseCheck =
       defn.mkList_contains(defn.mkList_apply(cases.map(p => c.literal(p._2))))(c.Expr(Ident(name.x1)))
+    val handlerName = name.asyncHander
 
     Block(List(
       // anonymous subclass of PartialFunction[Int, Unit]
       // TODO subclass AbstractPartialFunction
-      ClassDef(Modifiers(FINAL), name.anon, List(), Template(List(AppliedTypeTree(partFunIdent, List(intIdent, unitIdent))),
+      ClassDef(Modifiers(FINAL), handlerName, List(), Template(List(AppliedTypeTree(partFunIdent, List(intIdent, unitIdent))),
         emptyValDef, List(
           DefDef(Modifiers(), nme.CONSTRUCTOR, List(), List(List()), TypeTree(),
             Block(List(Apply(Select(Super(This(tpnme.EMPTY), tpnme.EMPTY), nme.CONSTRUCTOR), List())), c.literalUnit.tree)),
 
-          DefDef(Modifiers(), name.isDefinedAt, List(), List(List(ValDef(Modifiers(PARAM), name.x1, intIdent, EmptyTree))), TypeTree(),
+          DefDef(Modifiers(), name.isDefinedAt, List(), List(List(paramValDef(name.x1, definitions.IntClass))), TypeTree(),
             caseCheck.tree),
 
-          DefDef(Modifiers(), name.apply, List(), List(List(ValDef(Modifiers(PARAM), name.x1, intIdent, EmptyTree))), TypeTree(),
+          DefDef(Modifiers(), name.apply, List(), List(List(paramValDef(name.x1, definitions.IntClass))), TypeTree(),
             Match(Ident(name.x1), cases.map(_._1)) // combine all cases into a single match
           )
         ))
       )),
-      Apply(Select(New(Ident(name.anon)), nme.CONSTRUCTOR), List())
+      Apply(Select(New(Ident(handlerName)), nme.CONSTRUCTOR), List())
     )
   }
 
@@ -246,7 +253,7 @@ final class ExprBuilder[C <: Context, FS <: FutureSystem](val c: C, val futureSy
       val updateState = mkStateTree(nextState)
 
       val handlerTree =
-        Function(List(ValDef(Modifiers(PARAM), name.tr, TypeTree(tryType), EmptyTree)), Block(tryGetTree, updateState, mkResumeApply))
+        Function(List(paramValDef(name.tr, tryType)), Block(tryGetTree, updateState, mkResumeApply))
 
       futureSystemOps.onComplete(c.Expr(awaitable), c.Expr(handlerTree), execContext).tree
     }
@@ -552,6 +559,13 @@ final class ExprBuilder[C <: Context, FS <: FutureSystem](val c: C, val futureSy
       c.Expr(mkHandlerTreeFor(cases zip asyncStates.init.map(_.state)))
     }
 
+    def mkCombinedHandlerCases(): List[(CaseDef, Int)] = {
+      assert(asyncStates.size > 1)
+
+      val cases = for (state <- asyncStates.toList) yield state.mkHandlerCaseForState()
+      cases zip asyncStates.init.map(_.state)
+    }
+
     /* Builds the handler expression for a sequence of async states.
      */
     def mkHandlerExpr(): c.Expr[PartialFunction[Int, Unit]] = {