Merge pull request #176 from retronym/topic/extensions

Improve generated code and flexibility

Author: retronym

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

@@ -3,6 +3,7 @@
  */
 package scala.async.internal
 
+import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
 import language.existentials
 
@@ -117,16 +118,22 @@ trait ExprBuilder {
      *   <mkResumeApply>
      * }
      */
-    def ifIsFailureTree[T: WeakTypeTag](tryReference: => Tree) =
-      If(futureSystemOps.tryyIsFailure(c.Expr[futureSystem.Tryy[T]](tryReference)).tree,
-        Block(toList(futureSystemOps.completeProm[T](
-          c.Expr[futureSystem.Prom[T]](symLookup.memberRef(name.result)),
-          c.Expr[futureSystem.Tryy[T]](
-            TypeApply(Select(tryReference, newTermName("asInstanceOf")),
-              List(TypeTree(futureSystemOps.tryType[T]))))).tree),
-          Return(literalUnit)),
-        Block(List(tryGetTree(tryReference)), mkStateTree(nextState, symLookup))
-      )
+    def ifIsFailureTree[T: WeakTypeTag](tryReference: => Tree) = {
+      val getAndUpdateState = Block(List(tryGetTree(tryReference)), mkStateTree(nextState, symLookup))
+      if (asyncBase.futureSystem.emitTryCatch) {
+        If(futureSystemOps.tryyIsFailure(c.Expr[futureSystem.Tryy[T]](tryReference)).tree,
+          Block(toList(futureSystemOps.completeProm[T](
+            c.Expr[futureSystem.Prom[T]](symLookup.memberRef(name.result)),
+            c.Expr[futureSystem.Tryy[T]](
+              TypeApply(Select(tryReference, newTermName("asInstanceOf")),
+                List(TypeTree(futureSystemOps.tryType[T]))))).tree),
+            Return(literalUnit)),
+          getAndUpdateState
+        )
+      } else {
+        getAndUpdateState
+      }
+    }
 
     override def mkOnCompleteHandler[T: WeakTypeTag]: Option[CaseDef] = {
       Some(mkHandlerCase(onCompleteState, List(ifIsFailureTree[T](Ident(symLookup.applyTrParam)))))
@@ -401,9 +408,10 @@ trait ExprBuilder {
         val stateMemberSymbol = symLookup.stateMachineMember(name.state)
         val stateMemberRef = symLookup.memberRef(name.state)
         val body = Match(stateMemberRef, mkCombinedHandlerCases[T] ++ initStates.flatMap(_.mkOnCompleteHandler[T]) ++ List(CaseDef(Ident(nme.WILDCARD), EmptyTree, Throw(Apply(Select(New(Ident(defn.IllegalStateExceptionClass)), termNames.CONSTRUCTOR), List())))))
+        val body1 = eliminateDeadStates(body)
 
-        Try(
-          body,
+        maybeTry(
+          body1,
           List(
             CaseDef(
             Bind(name.t, Typed(Ident(nme.WILDCARD), Ident(defn.ThrowableClass))),
@@ -417,8 +425,67 @@ trait ExprBuilder {
               If(Apply(Ident(defn.NonFatalClass), List(Ident(name.t))), then, Throw(Ident(name.t)))
               then
             })), EmptyTree)
+      }
 
-        //body
+      // Identify dead states: `case <id> => { state = nextId; (); (); ... }, eliminated, and compact state ids to
+      // enable emission of a tableswitch.
+      private def eliminateDeadStates(m: Match): Tree = {
+        object DeadState {
+          private val liveStates = mutable.AnyRefMap[Integer, Integer]()
+          private val deadStates = mutable.AnyRefMap[Integer, Integer]()
+          private var compactedStateId = 1
+          for (CaseDef(Literal(Constant(stateId: Integer)), EmptyTree, body) <- m.cases) {
+            body match {
+              case _ if (stateId == 0) => liveStates(stateId) = stateId
+              case Block(Assign(_, Literal(Constant(nextState: Integer))) :: rest, expr) if (expr :: rest).forall(t => isLiteralUnit(t)) =>
+                deadStates(stateId) = nextState
+              case _ =>
+                liveStates(stateId) = compactedStateId
+                compactedStateId += 1
+            }
+          }
+          if (deadStates.nonEmpty)
+            AsyncUtils.vprintln(s"${deadStates.size} dead states eliminated")
+          def isDead(i: Integer) = deadStates.contains(i)
+          def translatedStateId(i: Integer, tree: Tree): Integer = {
+            def chaseDead(i: Integer): Integer = {
+              val replacement = deadStates.getOrNull(i)
+              if (replacement == null) i
+              else chaseDead(replacement)
+            }
+
+            val live = chaseDead(i)
+            liveStates.get(live) match {
+              case Some(x) => x
+              case None => sys.error(s"$live, $liveStates \n$deadStates\n$m\n\n====\n$tree")
+            }
+          }
+        }
+        val stateMemberSymbol = symLookup.stateMachineMember(name.state)
+        // - remove CaseDef-s for dead states
+        // - rewrite state transitions to dead states to instead transition to the
+        //   non-dead successor.
+        val elimDeadStateTransform = new Transformer {
+          override def transform(tree: Tree): Tree = tree match {
+            case as @ Assign(lhs, Literal(Constant(i: Integer))) if lhs.symbol == stateMemberSymbol =>
+              val replacement = DeadState.translatedStateId(i, as)
+              treeCopy.Assign(tree, lhs, Literal(Constant(replacement)))
+            case _: Match | _: CaseDef | _: Block | _: If =>
+              super.transform(tree)
+            case _ => tree
+          }
+        }
+        val cases1 = m.cases.flatMap {
+          case cd @ CaseDef(Literal(Constant(i: Integer)), EmptyTree, rhs) =>
+            if (DeadState.isDead(i)) Nil
+            else {
+              val replacement = DeadState.translatedStateId(i, cd)
+              val rhs1 = elimDeadStateTransform.transform(rhs)
+              treeCopy.CaseDef(cd, Literal(Constant(replacement)), EmptyTree, rhs1) :: Nil
+            }
+          case x => x :: Nil
+        }
+        treeCopy.Match(m, m.selector, cases1)
       }
 
       def forever(t: Tree): Tree = {

src/main/scala/scala/async/internal/FutureSystem.scala

@@ -74,6 +74,10 @@ trait FutureSystem {
   }
 
   def mkOps(c0: Context): Ops { val c: c0.type }
+
+  def freshenAllNames: Boolean = false
+  def emitTryCatch: Boolean = true
+  def resultFieldName: String = "result"
 }
 
 object ScalaConcurrentFutureSystem extends FutureSystem {

src/main/scala/scala/async/internal/TransformUtils.scala

@@ -17,28 +17,47 @@ private[async] trait TransformUtils {
   import c.internal._
   import decorators._
 
+  private object baseNames {
+
+    val matchRes = "matchres"
+    val ifRes = "ifres"
+    val bindSuffix = "$bind"
+    val completed = newTermName("completed")
+
+    val state = newTermName("state")
+    val result = newTermName(self.futureSystem.resultFieldName)
+    val execContext = newTermName("execContext")
+    val tr = newTermName("tr")
+    val t = newTermName("throwable")
+  }
+  
   object name {
-    val resume        = newTermName("resume")
-    val apply         = newTermName("apply")
-    val matchRes      = "matchres"
-    val ifRes         = "ifres"
-    val await         = "await"
-    val bindSuffix    = "$bind"
-    val completed     = newTermName("completed")
-
-    val state         = newTermName("state")
-    val result        = newTermName("result")
-    val execContext   = newTermName("execContext")
+    def matchRes      = maybeFresh(baseNames.matchRes)
+    def ifRes         = maybeFresh(baseNames.ifRes) 
+    def bindSuffix    = maybeFresh(baseNames.bindSuffix)
+    def completed     = maybeFresh(baseNames.completed) 
+
+    val state         = maybeFresh(baseNames.state)
+    val result        = baseNames.result
+    val execContext   = maybeFresh(baseNames.execContext)
+    val tr            = maybeFresh(baseNames.tr)
+    val t             = maybeFresh(baseNames.t)
+
+    val await = "await"
+    val resume = newTermName("resume")
+    val apply = newTermName("apply")
     val stateMachine  = newTermName(fresh("stateMachine"))
     val stateMachineT = stateMachine.toTypeName
-    val tr            = newTermName("tr")
-    val t             = newTermName("throwable")
 
+    def maybeFresh(name: TermName): TermName = if (self.asyncBase.futureSystem.freshenAllNames) fresh(name) else name
+    def maybeFresh(name: String): String = if (self.asyncBase.futureSystem.freshenAllNames) fresh(name) else name
     def fresh(name: TermName): TermName = c.freshName(name)
 
     def fresh(name: String): String = c.freshName(name)
   }
 
+  def maybeTry(block: Tree, catches: List[CaseDef], finalizer: Tree) = if (asyncBase.futureSystem.emitTryCatch) Try(block, catches, finalizer) else block
+
   def isAsync(fun: Tree) =
     fun.symbol == defn.Async_async