Rework dead state elimination to happen as part of expr builder

retronym · retronym · commit a8d4032f4250 · 2018-05-04T16:33:07.000+10:00
diff --git a/src/main/scala/scala/async/internal/AsyncTransform.scala b/src/main/scala/scala/async/internal/AsyncTransform.scala
@@ -130,10 +130,8 @@ trait AsyncTransform {
     AsyncUtils.vprintln(s"${c.macroApplication}")
     AsyncUtils.vprintln(s"ANF transform expands to:\n $anfTree")
     states foreach (s => AsyncUtils.vprintln(s))
-    AsyncUtils.vprintln("===== BEFORE DSE =====")
-    AsyncUtils.vprintln(block.toDot(afterDSE = false))
-    AsyncUtils.vprintln("===== AFTER DSE =====")
-    AsyncUtils.vprintln(block.toDot(afterDSE = true))
+    AsyncUtils.vprintln("===== DOT =====")
+    AsyncUtils.vprintln(block.toDot)
   }
 
   /**
diff --git a/src/main/scala/scala/async/internal/ExprBuilder.scala b/src/main/scala/scala/async/internal/ExprBuilder.scala
@@ -23,6 +23,8 @@ trait ExprBuilder {
   val labelDefStates = collection.mutable.Map[Symbol, Int]()
 
   trait AsyncState {
+    var switchId: Int = -1
+
     def state: Int
 
     def nextStates: Array[Int]
@@ -76,13 +78,13 @@ trait ExprBuilder {
       mkHandlerCase(state, stats)
 
     override val toString: String =
-      s"AsyncStateWithoutAwait #$state, nextStates = $nextStates"
+      s"AsyncStateWithoutAwait #$state, nextStates = ${nextStates.toList}"
   }
 
   /** A sequence of statements that concludes with an `await` call. The `onComplete`
     * handler will unconditionally transition to `nextState`.
     */
-  final class AsyncStateWithAwait(var stats: List[Tree], val state: Int, onCompleteState: Int, nextState: Int,
+  final class AsyncStateWithAwait(var stats: List[Tree], val state: Int, val onCompleteState: Int, nextState: Int,
                                   val awaitable: Awaitable, symLookup: SymLookup)
     extends AsyncState {
 
@@ -268,11 +270,11 @@ trait ExprBuilder {
     }
 
     // populate asyncStates
-    def add(stat: Tree): Unit = stat match {
+    def add(stat: Tree, afterState: Option[Int] = None): Unit = stat match {
       // the val name = await(..) pattern
       case vd @ ValDef(mods, name, tpt, Apply(fun, arg :: Nil)) if isAwait(fun) =>
         val onCompleteState = nextState()
-        val afterAwaitState = nextState()
+        val afterAwaitState = afterState.getOrElse(nextState())
         val awaitable = Awaitable(arg, stat.symbol, tpt.tpe, vd)
         asyncStates += stateBuilder.resultWithAwait(awaitable, onCompleteState, afterAwaitState) // complete with await
         currState = afterAwaitState
@@ -283,7 +285,7 @@ trait ExprBuilder {
 
         val thenStartState = nextState()
         val elseStartState = nextState()
-        val afterIfState = nextState()
+        val afterIfState = afterState.getOrElse(nextState())
 
         asyncStates +=
           // the two Int arguments are the start state of the then branch and the else branch, respectively
@@ -305,7 +307,7 @@ trait ExprBuilder {
         java.util.Arrays.setAll(caseStates, new IntUnaryOperator {
           override def applyAsInt(operand: Int): Int = nextState()
         })
-        val afterMatchState = nextState()
+        val afterMatchState = afterState.getOrElse(nextState())
 
         asyncStates +=
           stateBuilder.resultWithMatch(scrutinee, cases, caseStates, symLookup)
@@ -323,15 +325,16 @@ trait ExprBuilder {
         if containsAwait(rhs) || directlyAdjacentLabelDefs(ld).exists(containsAwait) =>
 
         val startLabelState = stateIdForLabel(ld.symbol)
-        val afterLabelState = nextState()
+        val afterLabelState = afterState.getOrElse(nextState())
         asyncStates += stateBuilder.resultWithLabel(startLabelState, symLookup)
         labelDefStates(ld.symbol) = startLabelState
         val builder = nestedBlockBuilder(rhs, startLabelState, afterLabelState)
         asyncStates ++= builder.asyncStates
         currState = afterLabelState
         stateBuilder = new AsyncStateBuilder(currState, symLookup)
       case b @ Block(stats, expr) =>
-        (stats :+ expr) foreach (add)
+        for (stat <- stats) add(stat)
+        add(expr, afterState = Some(endState))
       case _ =>
         checkForUnsupportedAwait(stat)
         stateBuilder += stat
@@ -346,7 +349,7 @@ trait ExprBuilder {
 
     def onCompleteHandler[T: WeakTypeTag]: Tree
 
-    def toDot(afterDSE: Boolean): String
+    def toDot: String
   }
 
   case class SymLookup(stateMachineClass: Symbol, applyTrParam: Symbol) {
@@ -371,11 +374,10 @@ trait ExprBuilder {
     val blockBuilder = new AsyncBlockBuilder(stats, expr, startState, endState, symLookup)
 
     new AsyncBlock {
-      val liveStates = mutable.AnyRefMap[Integer, Integer]()
-      val deadStates = mutable.AnyRefMap[Integer, Integer]()
+      val switchIds = mutable.AnyRefMap[Integer, Integer]()
 
       // render with http://graphviz.it/#/new
-      def toDot(afterDSE: Boolean): String = {
+      def toDot: String = {
         val states = asyncStates
         def toHtmlLabel(label: String, preText: String, builder: StringBuilder): Unit = {
           builder.append("<b>").append(label).append("</b>").append("<br/>")
@@ -390,39 +392,60 @@ trait ExprBuilder {
         val dotBuilder = new StringBuilder()
         dotBuilder.append("digraph {\n")
         def stateLabel(s: Int) = {
-          val beforeDseLabel = if (s == 0) "INITIAL" else if (s == Int.MaxValue) "TERMINAL" else if (s > 0) "S" + s else "C" + Math.abs(s)
-          if (afterDSE) {
-            "\"S" + liveStates.getOrElse(s, s) + " (" + beforeDseLabel + ")\""
-          } else {
-            beforeDseLabel
-          }
-
+          if (s == 0) "INITIAL" else if (s == Int.MaxValue) "TERMINAL" else switchIds.getOrElse(s, s).toString
         }
         val length = asyncStates.size
         for ((state, i) <- asyncStates.zipWithIndex) {
-          val liveStateIdOpt: Option[Int] = if (afterDSE) {
-            liveStates.get(state.state).map(_.intValue())
-          } else Some(state.state)
-          for (_ <- liveStateIdOpt) {
-            dotBuilder.append(s"""${stateLabel(state.state)} [label=""").append("<")
-            if (i != length - 1) {
-              val CaseDef(_, _, body) = state.mkHandlerCaseForState
-              toHtmlLabel(stateLabel(state.state), showCode(body), dotBuilder)
-            } else {
-              toHtmlLabel(stateLabel(state.state), state.allStats.map(showCode(_)).mkString("\n"), dotBuilder)
-            }
-            dotBuilder.append("> ]\n")
+          dotBuilder.append(s"""${stateLabel(state.state)} [label=""").append("<")
+          if (i != length - 1) {
+            val CaseDef(_, _, body) = state.mkHandlerCaseForState
+            toHtmlLabel(stateLabel(state.state), showCode(body), dotBuilder)
+          } else {
+            toHtmlLabel(stateLabel(state.state), state.allStats.map(showCode(_)).mkString("\n"), dotBuilder)
           }
+          dotBuilder.append("> ]\n")
         }
-        for (state <- states; if liveStates.contains(state.state); succ <- state.nextStates) {
+        for (state <- states; succ <- state.nextStates) {
           dotBuilder.append(s"""${stateLabel(state.state)} -> ${stateLabel(succ)}""")
           dotBuilder.append("\n")
         }
         dotBuilder.append("}\n")
         dotBuilder.toString
       }
 
-      def asyncStates = blockBuilder.asyncStates.toList
+      lazy val asyncStates: List[AsyncState] = filterStates
+
+      def filterStates = {
+        val all = blockBuilder.asyncStates.toList
+        val (initial :: rest) = all
+        val map = all.iterator.map(x => (x.state, x)).toMap
+        var seen = mutable.HashSet[Int]()
+        def loop(state: AsyncState): Unit = {
+          seen.add(state.state)
+          for (i <- state.nextStates) {
+            if (i != Int.MaxValue && !seen.contains(i)) {
+              loop(map(i))
+            }
+          }
+        }
+        loop(initial)
+        val live = rest.filter(state => seen(state.state))
+        var nextSwitchId = 1
+        (initial :: live).foreach { state =>
+          val switchId = nextSwitchId
+          switchIds(state.state) = switchId
+          nextSwitchId += 1
+          state match {
+            case state: AsyncStateWithAwait =>
+              val switchId = nextSwitchId
+              switchIds(state.onCompleteState) = switchId
+              nextSwitchId += 1
+            case _ =>
+          }
+        }
+        initial :: live
+
+      }
 
       def mkCombinedHandlerCases[T: WeakTypeTag]: List[CaseDef] = {
         val caseForLastState: CaseDef = {
@@ -488,43 +511,14 @@ trait ExprBuilder {
       // 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 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)
+              val replacement = switchIds(i)
               treeCopy.Assign(tree, lhs, Literal(Constant(replacement)))
             case _: Match | _: CaseDef | _: Block | _: If =>
               super.transform(tree)
@@ -533,12 +527,9 @@ trait ExprBuilder {
         }
         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
-            }
+            val replacement = switchIds(i)
+            val rhs1 = elimDeadStateTransform.transform(rhs)
+            treeCopy.CaseDef(cd, Literal(Constant(replacement)), EmptyTree, rhs1) :: Nil
           case x => x :: Nil
         }
         treeCopy.Match(m, m.selector, cases1)
diff --git a/src/main/scala/scala/async/internal/TransformUtils.scala b/src/main/scala/scala/async/internal/TransformUtils.scala
@@ -39,6 +39,11 @@ private[async] trait TransformUtils {
       true
     case _ => false
   }
+  def isNullAssignment(t: Tree) = t match {
+    case Assign(Ident(_), Literal(Constant(null))) =>
+      true
+    case _ => false
+  }
 
   def isPastTyper =
     c.universe.asInstanceOf[scala.reflect.internal.SymbolTable].isPastTyper
