Avoid some boxing of state ids during transform

retronym · retronym · commit 5fb8657cbebc · 2018-06-07T15:14:34.000+10:00
diff --git a/src/main/scala/scala/async/internal/ExprBuilder.scala b/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,10 @@ 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, new IntUnaryOperator {
+          override def applyAsInt(operand: Int): Int = nextState()
+        })
         val afterMatchState = nextState()
 
         asyncStates +=
diff --git a/src/main/scala/scala/async/internal/LiveVariables.scala b/src/main/scala/scala/async/internal/LiveVariables.scala
@@ -1,5 +1,10 @@
 package scala.async.internal
 
+import java.util
+import java.util.function.{IntConsumer, IntPredicate}
+
+import scala.collection.immutable.IntMap
+
 trait LiveVariables {
   self: AsyncMacro =>
   import c.universe._
@@ -17,19 +22,22 @@ trait LiveVariables {
   def fieldsToNullOut(asyncStates: List[AsyncState], liftables: List[Tree]): Map[Int, List[Tree]] = {
     // live variables analysis:
     // the result map indicates in which states a given field should be nulled out
-    val liveVarsMap: Map[Tree, Set[Int]] = liveVars(asyncStates, liftables)
+    val liveVarsMap: Map[Tree, StateSet] = liveVars(asyncStates, liftables)
 
     var assignsOf = Map[Int, List[Tree]]()
 
-    for ((fld, where) <- liveVarsMap; state <- where)
-      assignsOf get state match {
-        case None =>
-          assignsOf += (state -> List(fld))
-        case Some(trees) if !trees.exists(_.symbol == fld.symbol) =>
-          assignsOf += (state -> (fld +: trees))
-        case _ =>
-          /* do nothing */
-      }
+    for ((fld, where) <- liveVarsMap) {
+      where.foreach { new IntConsumer { def accept(state: Int): Unit = {
+        assignsOf get state match {
+          case None =>
+            assignsOf += (state -> List(fld))
+          case Some(trees) if !trees.exists(_.symbol == fld.symbol) =>
+            assignsOf += (state -> (fld +: trees))
+          case _ =>
+            // do nothing
+        }
+      }}}
+    }
 
     assignsOf
   }
@@ -46,9 +54,9 @@ trait LiveVariables {
    *  @param   liftables   the lifted fields
    *  @return              a map which indicates for a given field (the key) the states in which it should be nulled out
    */
-  def liveVars(asyncStates: List[AsyncState], liftables: List[Tree]): Map[Tree, Set[Int]] = {
+  def liveVars(asyncStates: List[AsyncState], liftables: List[Tree]): Map[Tree, StateSet] = {
     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 +130,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 +182,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 +192,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 +209,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
         }
       }
@@ -210,53 +236,64 @@ trait LiveVariables {
       }
     }
 
-    def lastUsagesOf(field: Tree, at: AsyncState): Set[Int] = {
+    def lastUsagesOf(field: Tree, at: AsyncState): StateSet = {
       val avoid = scala.collection.mutable.HashSet[AsyncState]()
 
-      def lastUsagesOf0(field: Tree, at: AsyncState): Set[Int] = {
-        if (avoid(at)) Set()
+      val result = new StateSet
+      def lastUsagesOf0(field: Tree, at: AsyncState): Unit = {
+        if (avoid(at)) ()
         else if (captured(field.symbol)) {
-          Set()
+          ()
         }
         else LVentry get at.state match {
           case Some(fields) if fields.contains(field.symbol) =>
-            Set(at.state)
+            result += at.state
           case _ =>
             avoid += at
-            val preds = asyncStates.filter(_.nextStates.contains(at.state)).toSet
-            preds.flatMap(p => lastUsagesOf0(field, p))
+            for (state <- asyncStates) {
+              if (contains(state.nextStates, at.state)) {
+                lastUsagesOf0(field, state)
+              }
+            }
         }
       }
 
       lastUsagesOf0(field, at)
+      result
     }
 
-    val lastUsages: Map[Tree, Set[Int]] =
-      liftables.map(fld => fld -> lastUsagesOf(fld, finalState)).toMap
+    val lastUsages: Map[Tree, StateSet] =
+      liftables.iterator.map(fld => fld -> lastUsagesOf(fld, finalState)).toMap
 
     if(AsyncUtils.verbose) {
       for ((fld, lastStates) <- lastUsages)
-        AsyncUtils.vprintln(s"field ${fld.symbol.name} is last used in states ${lastStates.mkString(", ")}")
+        AsyncUtils.vprintln(s"field ${fld.symbol.name} is last used in states ${lastStates.iterator.mkString(", ")}")
     }
 
-    val nullOutAt: Map[Tree, Set[Int]] =
+    val nullOutAt: Map[Tree, StateSet] =
       for ((fld, lastStates) <- lastUsages) yield {
-        val killAt = lastStates.flatMap { s =>
-          if (s == finalState.state) Set()
-          else {
+        var result = new StateSet
+        lastStates.foreach(new IntConsumer { def accept(s: Int): Unit = {
+          if (s != finalState.state) {
             val lastAsyncState = asyncStates.find(_.state == s).get
             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))
+            var i = 0
+            while (i < succNums.length) {
+              val num = succNums(i)
+              if (!isPred(num, s) && !LVentry(num).contains(fld.symbol))
+                result += num
+              i += 1
+            }
           }
-        }
-        (fld, killAt)
+        }})
+        (fld, result)
       }
 
     if(AsyncUtils.verbose) {
       for ((fld, killAt) <- nullOutAt)
-        AsyncUtils.vprintln(s"field ${fld.symbol.name} should be nulled out in states ${killAt.mkString(", ")}")
+        AsyncUtils.vprintln(s"field ${fld.symbol.name} should be nulled out in states ${killAt.iterator.mkString(", ")}")
     }
 
     nullOutAt
diff --git a/src/main/scala/scala/async/internal/StateSet.scala b/src/main/scala/scala/async/internal/StateSet.scala
@@ -0,0 +1,27 @@
+/*
+ * Copyright (C) 2018 Lightbend Inc. <http://www.lightbend.com>
+ */
+package scala.async.internal
+
+import java.util
+import java.util.function.{Consumer, IntConsumer}
+
+import scala.collection.JavaConverters.{asScalaIteratorConverter, iterableAsScalaIterableConverter}
+import scala.collection.mutable
+
+// Set for StateIds, which are either small positive integers or -symbolID.
+final class StateSet {
+  private var bitSet = new java.util.BitSet()
+  private var caseSet = new util.HashSet[Integer]()
+  def +=(stateId: Int): Unit = if (stateId > 0) bitSet.set(stateId) else caseSet.add(stateId)
+  def contains(stateId: Int): Boolean = if (stateId > 0 && stateId < 1024) bitSet.get(stateId) else caseSet.contains(stateId)
+  def iterator: Iterator[Integer] = {
+    bitSet.stream().iterator().asScala ++ caseSet.asScala.iterator
+  }
+  def foreach(f: IntConsumer): Unit = {
+    bitSet.stream().forEach(f)
+    caseSet.stream().forEach(new Consumer[Integer] {
+      override def accept(value: Integer): Unit = f.accept(value)
+    })
+  }
+}
