Avoid some boxing of state ids during transform

Author: retronym

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

@@ -3,6 +3,8 @@
  */
 package scala.async.internal
 
+import java.util.function.IntUnaryOperator
+
 import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
 import language.existentials
@@ -23,7 +25,7 @@ trait ExprBuilder {
   trait AsyncState {
     def state: Int
 
-    def nextStates: List[Int]
+    def nextStates: Array[Int]
 
     def mkHandlerCaseForState[T: WeakTypeTag]: CaseDef
 
@@ -55,8 +57,8 @@ trait ExprBuilder {
   final class SimpleAsyncState(var stats: List[Tree], val state: Int, nextState: Int, symLookup: SymLookup)
     extends AsyncState {
 
-    def nextStates: List[Int] =
-      List(nextState)
+    val nextStates: Array[Int] =
+      Array(nextState)
 
     def mkHandlerCaseForState[T: WeakTypeTag]: CaseDef = {
       mkHandlerCase(state, treesThenStats(mkStateTree(nextState, symLookup) :: Nil))
@@ -69,7 +71,7 @@ trait ExprBuilder {
   /** A sequence of statements with a conditional transition to the next state, which will represent
     * a branch of an `if` or a `match`.
     */
-  final class AsyncStateWithoutAwait(var stats: List[Tree], val state: Int, val nextStates: List[Int]) extends AsyncState {
+  final class AsyncStateWithoutAwait(var stats: List[Tree], val state: Int, val nextStates: Array[Int]) extends AsyncState {
     override def mkHandlerCaseForState[T: WeakTypeTag]: CaseDef =
       mkHandlerCase(state, stats)
 
@@ -84,8 +86,8 @@ trait ExprBuilder {
                                   val awaitable: Awaitable, symLookup: SymLookup)
     extends AsyncState {
 
-    def nextStates: List[Int] =
-      List(nextState)
+    val nextStates: Array[Int] =
+      Array(nextState)
 
     override def mkHandlerCaseForState[T: WeakTypeTag]: CaseDef = {
       val fun = This(tpnme.EMPTY)
@@ -191,7 +193,7 @@ trait ExprBuilder {
     def resultWithIf(condTree: Tree, thenState: Int, elseState: Int): AsyncState = {
       def mkBranch(state: Int) = mkStateTree(state, symLookup)
       this += If(condTree, mkBranch(thenState), mkBranch(elseState))
-      new AsyncStateWithoutAwait(stats.toList, state, List(thenState, elseState))
+      new AsyncStateWithoutAwait(stats.toList, state, Array(thenState, elseState))
     }
 
     /**
@@ -204,7 +206,7 @@ trait ExprBuilder {
      * @param caseStates      starting state of the right-hand side of the each case
      * @return                an `AsyncState` representing the match expression
      */
-    def resultWithMatch(scrutTree: Tree, cases: List[CaseDef], caseStates: List[Int], symLookup: SymLookup): AsyncState = {
+    def resultWithMatch(scrutTree: Tree, cases: List[CaseDef], caseStates: Array[Int], symLookup: SymLookup): AsyncState = {
       // 1. build list of changed cases
       val newCases = for ((cas, num) <- cases.zipWithIndex) yield cas match {
         case CaseDef(pat, guard, rhs) =>
@@ -218,7 +220,7 @@ trait ExprBuilder {
 
     def resultWithLabel(startLabelState: Int, symLookup: SymLookup): AsyncState = {
       this += mkStateTree(startLabelState, symLookup)
-      new AsyncStateWithoutAwait(stats.toList, state, List(startLabelState))
+      new AsyncStateWithoutAwait(stats.toList, state, Array(startLabelState))
     }
 
     override def toString: String = {
@@ -299,7 +301,8 @@ trait ExprBuilder {
       case Match(scrutinee, cases) if containsAwait(stat) =>
         checkForUnsupportedAwait(scrutinee)
 
-        val caseStates = cases.map(_ => nextState())
+        val caseStates = new Array[Int](cases.length)
+        java.util.Arrays.setAll(caseStates, (_ => nextState()): IntUnaryOperator)
         val afterMatchState = nextState()
 
         asyncStates +=

src/main/scala/scala/async/internal/LiveVariables.scala

@@ -1,5 +1,7 @@
 package scala.async.internal
 
+import scala.collection.immutable.IntMap
+
 trait LiveVariables {
   self: AsyncMacro =>
   import c.universe._
@@ -48,7 +50,7 @@ trait LiveVariables {
    */
   def liveVars(asyncStates: List[AsyncState], liftables: List[Tree]): Map[Tree, Set[Int]] = {
     val liftedSyms: Set[Symbol] = // include only vars
-      liftables.filter {
+      liftables.iterator.filter {
         case ValDef(mods, _, _, _) => mods.hasFlag(MUTABLE)
         case _ => false
       }.map(_.symbol).toSet
@@ -122,20 +124,30 @@ trait LiveVariables {
      * A state `i` is contained in the list that is the value to which
      * key `j` maps iff control can flow from state `j` to state `i`.
      */
-    val cfg: Map[Int, List[Int]] = asyncStates.map(as => as.state -> as.nextStates).toMap
+    val cfg: Map[Int, Array[Int]] = {
+      var res = IntMap.empty[Array[Int]]
+
+      for (as <- asyncStates) res = res.updated(as.state, as.nextStates)
+      res
+    }
 
     /** Tests if `state1` is a predecessor of `state2`.
      */
     def isPred(state1: Int, state2: Int): Boolean = {
-      val seen = scala.collection.mutable.HashSet[Int]()
+      val seen = new StateSet()
 
       def isPred0(state1: Int, state2: Int): Boolean = 
         if(state1 == state2) false
-        else if (seen(state1)) false  // breaks cycles in the CFG
+        else if (seen.contains(state1)) false  // breaks cycles in the CFG
         else cfg get state1 match {
           case Some(nextStates) =>
             seen += state1
-            nextStates.contains(state2) || nextStates.exists(isPred0(_, state2))
+            var i = 0
+            while (i < nextStates.length) {
+              if (nextStates(i) == state2 || isPred0(nextStates(i), state2)) return true
+              i += 1
+            }
+            false
           case None =>
             false
         }
@@ -164,8 +176,8 @@ trait LiveVariables {
      * 7. repeat if something has changed
      */
 
-    var LVentry = Map[Int, Set[Symbol]]() withDefaultValue Set[Symbol]()
-    var LVexit  = Map[Int, Set[Symbol]]() withDefaultValue Set[Symbol]()
+    var LVentry = IntMap[Set[Symbol]]() withDefaultValue Set[Symbol]()
+    var LVexit  = IntMap[Set[Symbol]]() withDefaultValue Set[Symbol]()
 
     // All fields are declared to be dead at the exit of the final async state, except for the ones
     // that cannot be nulled out at all (those in noNull), because they have been captured by a nested def.
@@ -174,6 +186,14 @@ trait LiveVariables {
     var currStates = List(finalState)    // start at final state
     var captured: Set[Symbol] = Set()
 
+    def contains(as: Array[Int], a: Int): Boolean = {
+      var i = 0
+      while (i < as.length) {
+        if (as(i) == a) return true
+        i += 1
+      }
+      false
+    }
     while (!currStates.isEmpty) {
       var entryChanged: List[AsyncState] = Nil
 
@@ -183,19 +203,19 @@ trait LiveVariables {
         captured ++= referenced.captured
         val LVentryNew = LVexit(cs.state) ++ referenced.used
         if (!LVentryNew.sameElements(LVentryOld)) {
-          LVentry = LVentry + (cs.state -> LVentryNew)
+          LVentry = LVentry.updated(cs.state, LVentryNew)
           entryChanged ::= cs
         }
       }
 
-      val pred = entryChanged.flatMap(cs => asyncStates.filter(_.nextStates.contains(cs.state)))
+      val pred = entryChanged.flatMap(cs => asyncStates.filter(state => contains(state.nextStates, cs.state)))
       var exitChanged: List[AsyncState] = Nil
 
       for (p <- pred) {
         val LVexitOld = LVexit(p.state)
         val LVexitNew = p.nextStates.flatMap(succ => LVentry(succ)).toSet
         if (!LVexitNew.sameElements(LVexitOld)) {
-          LVexit = LVexit + (p.state -> LVexitNew)
+          LVexit = LVexit.updated(p.state, LVexitNew)
           exitChanged ::= p
         }
       }
@@ -223,7 +243,7 @@ trait LiveVariables {
             Set(at.state)
           case _ =>
             avoid += at
-            val preds = asyncStates.filter(_.nextStates.contains(at.state)).toSet
+            val preds = asyncStates.filter(state => contains(state.nextStates, at.state)).toSet
             preds.flatMap(p => lastUsagesOf0(field, p))
         }
       }
@@ -248,7 +268,7 @@ trait LiveVariables {
             val succNums       = lastAsyncState.nextStates
             // all successor states that are not indirect predecessors
             // filter out successor states where the field is live at the entry
-            succNums.filter(num => !isPred(num, s)).filterNot(num => LVentry(num).contains(fld.symbol))
+            succNums.iterator.filter(num => !isPred(num, s)).filterNot(num => LVentry(num).contains(fld.symbol))
           }
         }
         (fld, killAt)

src/main/scala/scala/async/internal/StateSet.scala

@@ -0,0 +1,14 @@
+package scala.async.internal
+
+import java.util
+
+import scala.collection.mutable
+
+// Set for StateIds, which are either small positive integers or -symbolID.
+final class StateSet {
+  private var bitSet = mutable.BitSet()
+  private var caseSet = new util.HashSet[Integer]()
+  def +=(stateId: Int): Unit = if (stateId > 0) bitSet += stateId else caseSet.add(stateId)
+  def contains(stateId: Int): Boolean = if (stateId > 0) bitSet.contains(stateId) else caseSet.contains(stateId)
+
+}