Less boxing in field nulling analysis

Author: retronym

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

@@ -1,5 +1,7 @@
 package scala.async.internal
 
+import java.util
+
 import scala.collection.immutable.IntMap
 
 trait LiveVariables {
@@ -19,19 +21,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 { (state: Int) =>
+        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
   }
@@ -48,7 +53,7 @@ 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.iterator.filter {
         case ValDef(mods, _, _, _) => mods.hasFlag(MUTABLE)
@@ -230,53 +235,58 @@ 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(state => contains(state.nextStates, 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 { s =>
+          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.iterator.filter(num => !isPred(num, s)).filterNot(num => LVentry(num).contains(fld.symbol))
+            util.Arrays.stream(succNums).filter(num => !isPred(num, s)).filter(num => !LVentry(num).contains(fld.symbol)).forEach(result += _)
           }
         }
-        (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

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

@@ -1,14 +1,25 @@
+/*
+ * Copyright (C) 2018 Lightbend Inc. <http://www.lightbend.com>
+ */
 package scala.async.internal
 
 import java.util
+import java.util.function.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 = mutable.BitSet()
+  private var bitSet = new java.util.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)
-
+  def +=(stateId: Int): Unit = if (stateId > 0) bitSet.set(stateId) else caseSet.add(stateId)
+  def contains(stateId: Int): Boolean = if (stateId > 0) 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(integer => f.accept(integer))
+  }
 }