Group call graph by method

Change the hash maps in the call graph to index callsites and
closure instantiations by the containing method. This is beneficial
for implementing inliner heuristics and the closure optimizer.

3 files changed, 100 insertions, 69 deletions

diff --git a/src/compiler/scala/tools/nsc/backend/jvm/opt/CallGraph.scala b/src/compiler/scala/tools/nsc/backend/jvm/opt/CallGraph.scala
index 96455c0e38..e11ae08a8c 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/CallGraph.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/CallGraph.scala
@@ -11,7 +11,7 @@ import scala.reflect.internal.util.{NoPosition, Position}
 import scala.tools.asm.tree.analysis.{Value, Analyzer, BasicInterpreter}
 import scala.tools.asm.{Opcodes, Type, Handle}
 import scala.tools.asm.tree._
-import scala.collection.concurrent
+import scala.collection.{concurrent, mutable}
 import scala.collection.convert.decorateAsScala._
 import scala.tools.nsc.backend.jvm.BTypes.InternalName
 import scala.tools.nsc.backend.jvm.BackendReporting._
@@ -22,25 +22,41 @@ import BytecodeUtils._
 class CallGraph[BT <: BTypes](val btypes: BT) {
   import btypes._
 
-  val callsites: concurrent.Map[MethodInsnNode, Callsite] = recordPerRunCache(concurrent.TrieMap.empty)
+  /**
+   * The call graph contains the callsites in the program being compiled.
+   *
+   * Indexing the call graph by the containing MethodNode and the invocation MethodInsnNode allows
+   * finding callsites efficiently. For example, an inlining heuristic might want to know all
+   * callsites withing a callee method.
+   *
+   * Note that call graph entries (Callsite instances) keep a reference to the invocation
+   * MethodInsnNode, which keeps all AbstractInsnNodes of the method reachable. Adding classes
+   * from the classpath to the call graph (in addition to classes being compiled) may prevent
+   * method instruction nodes from being GCd. The ByteCodeRepository has a fixed size cache for
+   * parsed ClassNodes - keeping all ClassNodes alive consumed too much memory.
+   * The call graph is less problematic because only methods being called are kept alive, not entire
+   * classes. But we should keep an eye on this.
+   */
+  val callsites: mutable.Map[MethodNode, Map[MethodInsnNode, Callsite]] = recordPerRunCache(concurrent.TrieMap.empty withDefaultValue Map.empty)
 
-  val closureInstantiations: concurrent.Map[InvokeDynamicInsnNode, ClosureInstantiation] = recordPerRunCache(concurrent.TrieMap.empty)
+  /**
+   * Closure instantiations in the program being compiled.
+   *
+   * Indexing closure instantiations by the containing MethodNode is beneficial for the closure
+   * optimizer: finding callsites to re-write requires running a producers-consumers analysis on
+   * the method. Here the closure instantiations are already grouped by method.
+   */
+  val closureInstantiations: mutable.Map[MethodNode, Map[InvokeDynamicInsnNode, ClosureInstantiation]] = recordPerRunCache(concurrent.TrieMap.empty withDefaultValue Map.empty)
 
   def addClass(classNode: ClassNode): Unit = {
     val classType = classBTypeFromClassNode(classNode)
-    for {
-      m <- classNode.methods.asScala
-      (calls, closureInits) = analyzeCallsites(m, classType)
-    } {
-      calls foreach (callsite => callsites(callsite.callsiteInstruction) = callsite)
-      closureInits foreach (lmf => closureInstantiations(lmf.indy) = ClosureInstantiation(lmf, m, classType))
-    }
+    classNode.methods.asScala.foreach(addMethod(_, classType))
   }
 
   /**
    * Returns a list of callsites in the method, plus a list of closure instantiation indy instructions.
    */
-  def analyzeCallsites(methodNode: MethodNode, definingClass: ClassBType): (List[Callsite], List[LambdaMetaFactoryCall]) = {
+  def addMethod(methodNode: MethodNode, definingClass: ClassBType): Unit = {
 
     case class CallsiteInfo(safeToInline: Boolean, safeToRewrite: Boolean,
                             annotatedInline: Boolean, annotatedNoInline: Boolean,
@@ -128,8 +144,8 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
       case _ => false
     }
 
-    val callsites = new collection.mutable.ListBuffer[Callsite]
-    val closureInstantiations = new collection.mutable.ListBuffer[LambdaMetaFactoryCall]
+    var methodCallsites = Map.empty[MethodInsnNode, Callsite]
+    var methodClosureInstantiations = Map.empty[InvokeDynamicInsnNode, ClosureInstantiation]
 
     methodNode.instructions.iterator.asScala foreach {
       case call: MethodInsnNode =>
@@ -162,7 +178,7 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
           receiverNotNullByAnalysis(call, numArgs)
         }
 
-        callsites += Callsite(
+        methodCallsites += call -> Callsite(
           callsiteInstruction = call,
           callsiteMethod = methodNode,
           callsiteClass = definingClass,
@@ -174,15 +190,45 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
         )
 
       case LambdaMetaFactoryCall(indy, samMethodType, implMethod, instantiatedMethodType) =>
-        closureInstantiations += LambdaMetaFactoryCall(indy, samMethodType, implMethod, instantiatedMethodType)
+        methodClosureInstantiations += indy -> ClosureInstantiation(
+          LambdaMetaFactoryCall(indy, samMethodType, implMethod, instantiatedMethodType),
+          methodNode,
+          definingClass)
 
       case _ =>
     }
 
-    (callsites.toList, closureInstantiations.toList)
+    callsites(methodNode) = methodCallsites
+    closureInstantiations(methodNode) = methodClosureInstantiations
   }
 
-  /**
+  def removeCallsite(invocation: MethodInsnNode, methodNode: MethodNode): Option[Callsite] = {
+    val methodCallsites = callsites(methodNode)
+    val newCallsites = methodCallsites - invocation
+    if (newCallsites.isEmpty) callsites.remove(methodNode)
+    else callsites(methodNode) = newCallsites
+    methodCallsites.get(invocation)
+  }
+
+  def addCallsite(callsite: Callsite): Unit = {
+    val methodCallsites = callsites(callsite.callsiteMethod)
+    callsites(callsite.callsiteMethod) = methodCallsites + (callsite.callsiteInstruction -> callsite)
+  }
+
+  def removeClosureInstantiation(indy: InvokeDynamicInsnNode, methodNode: MethodNode): Option[ClosureInstantiation] = {
+    val methodClosureInits = closureInstantiations(methodNode)
+    val newClosureInits = methodClosureInits - indy
+    if (newClosureInits.isEmpty) closureInstantiations.remove(methodNode)
+    else closureInstantiations(methodNode) = newClosureInits
+    methodClosureInits.get(indy)
+  }
+
+  def addClosureInstantiation(closureInit: ClosureInstantiation) = {
+    val methodClosureInits = closureInstantiations(closureInit.ownerMethod)
+    closureInstantiations(closureInit.ownerMethod) = methodClosureInits + (closureInit.lambdaMetaFactoryCall.indy -> closureInit)
+  }
+
+    /**
    * A callsite in the call graph.
    *
    * @param callsiteInstruction The invocation instruction
diff --git a/src/compiler/scala/tools/nsc/backend/jvm/opt/ClosureOptimizer.scala b/src/compiler/scala/tools/nsc/backend/jvm/opt/ClosureOptimizer.scala
index 92b9b34006..3fbc374a93 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/ClosureOptimizer.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/ClosureOptimizer.scala
@@ -70,23 +70,16 @@ class ClosureOptimizer[BT <: BTypes](val btypes: BT) {
       }
     }
 
-    // Grouping the closure instantiations by method allows running the ProdConsAnalyzer only once per
-    // method. Also sort the instantiations: If there are multiple closure instantiations in a method,
-    // closure invocations need to be re-written in a consistent order for bytecode stability. The local
-    // variable slots for storing captured values depends on the order of rewriting.
-    val closureInstantiationsByMethod: Map[MethodNode, immutable.TreeSet[ClosureInstantiation]] = {
-      closureInstantiations.values.groupBy(_.ownerMethod).mapValues(immutable.TreeSet.empty ++ _)
-    }
-
     // For each closure instantiation, a list of callsites of the closure that can be re-written
     // If a callsite cannot be rewritten, for example because the lambda body method is not accessible,
     // a warning is returned instead.
     val callsitesToRewrite: List[(ClosureInstantiation, List[Either[RewriteClosureApplyToClosureBodyFailed, (MethodInsnNode, Int)]])] = {
-      closureInstantiationsByMethod.iterator.flatMap({
+      closureInstantiations.iterator.flatMap({
         case (methodNode, closureInits) =>
           // A lazy val to ensure the analysis only runs if necessary (the value is passed by name to `closureCallsites`)
-          lazy val prodCons = new ProdConsAnalyzer(methodNode, closureInits.head.ownerClass.internalName)
-          closureInits.iterator.map(init => (init, closureCallsites(init, prodCons)))
+          lazy val prodCons = new ProdConsAnalyzer(methodNode, closureInits.valuesIterator.next.ownerClass.internalName)
+          val sortedInits = immutable.TreeSet.empty ++ closureInits.values
+          sortedInits.iterator.map(init => (init, closureCallsites(init, prodCons)))
       }).toList // mapping to a list (not a map) to keep the sorting of closureInstantiationsByMethod
     }
 
@@ -162,7 +155,7 @@ class ClosureOptimizer[BT <: BTypes](val btypes: BT) {
           isAccessible
         }
 
-        def pos = callGraph.callsites.get(invocation).map(_.callsitePosition).getOrElse(NoPosition)
+        def pos = callGraph.callsites(ownerMethod).get(invocation).map(_.callsitePosition).getOrElse(NoPosition)
         val stackSize: Either[RewriteClosureApplyToClosureBodyFailed, Int] = bodyAccessible match {
           case Left(w)      => Left(RewriteClosureAccessCheckFailed(pos, w))
           case Right(false) => Left(RewriteClosureIllegalAccess(pos, ownerClass.internalName))
@@ -237,7 +230,7 @@ class ClosureOptimizer[BT <: BTypes](val btypes: BT) {
     ownerMethod.instructions.remove(invocation)
 
     // update the call graph
-    val originalCallsite = callGraph.callsites.remove(invocation)
+    val originalCallsite = callGraph.removeCallsite(invocation, ownerMethod)
 
     // the method node is needed for building the call graph entry
     val bodyMethod = byteCodeRepository.methodNode(lambdaBodyHandle.getOwner, lambdaBodyHandle.getName, lambdaBodyHandle.getDesc)
@@ -266,7 +259,7 @@ class ClosureOptimizer[BT <: BTypes](val btypes: BT) {
     //     (corresponding to the receiver) must be non-null"
     // Explanation: If the lambda body method is non-static, the receiver is a captured
     // value. It can only be captured within some instance method, so we know it's non-null.
-    callGraph.callsites(bodyInvocation) = bodyMethodCallsite
+    callGraph.addCallsite(bodyMethodCallsite)
   }
 
   /**
diff --git a/src/compiler/scala/tools/nsc/backend/jvm/opt/Inliner.scala b/src/compiler/scala/tools/nsc/backend/jvm/opt/Inliner.scala
index 8477f5461a..c5d2d212f4 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/Inliner.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/Inliner.scala
@@ -27,8 +27,8 @@ class Inliner[BT <: BTypes](val btypes: BT) {
 
   def eliminateUnreachableCodeAndUpdateCallGraph(methodNode: MethodNode, definingClass: InternalName): Unit = {
     localOpt.minimalRemoveUnreachableCode(methodNode, definingClass) foreach {
-      case invocation: MethodInsnNode  => callGraph.callsites.remove(invocation)
-      case indy: InvokeDynamicInsnNode => callGraph.closureInstantiations.remove(indy)
+      case invocation: MethodInsnNode  => callGraph.removeCallsite(invocation, methodNode)
+      case indy: InvokeDynamicInsnNode => callGraph.removeClosureInstantiation(indy, methodNode)
       case _ =>
     }
   }
@@ -48,7 +48,7 @@ class Inliner[BT <: BTypes](val btypes: BT) {
 
       // DCE above removes unreachable callsites from the call graph. If the inlining request denotes
       // such an eliminated callsite, do nothing.
-      if (callGraph.callsites contains request.callsiteInstruction) {
+      if (callGraph.callsites(request.callsiteMethod).contains(request.callsiteInstruction)) {
         val r = inline(request.callsiteInstruction, request.callsiteStackHeight, request.callsiteMethod, request.callsiteClass,
           callee.callee, callee.calleeDeclarationClass,
           request.receiverKnownNotNull, keepLineNumbers = false)
@@ -90,7 +90,7 @@ class Inliner[BT <: BTypes](val btypes: BT) {
    * requests is allowed to have cycles, and the callsites can appear in any order.
    */
   def selectCallsitesForInlining: List[Callsite] = {
-    callsites.valuesIterator.filter({
+    callsites.valuesIterator.flatMap(_.valuesIterator.filter({
       case callsite @ Callsite(_, _, _, Right(Callee(callee, calleeDeclClass, safeToInline, _, annotatedInline, _, warning)), _, _, _, pos) =>
         val res = doInlineCallsite(callsite)
 
@@ -118,7 +118,7 @@ class Inliner[BT <: BTypes](val btypes: BT) {
         if (warning.emitWarning(compilerSettings))
           backendReporting.inlinerWarning(pos, s"failed to determine if ${ins.name} should be inlined:\n$warning")
         false
-    }).toList
+    })).toList
   }
 
   /**
@@ -136,7 +136,7 @@ class Inliner[BT <: BTypes](val btypes: BT) {
     // Rewriting final trait method callsites to the implementation class enables inlining.
     // We cannot just iterate over the values of the `callsites` map because the rewrite changes the
     // map. Therefore we first copy the values to a list.
-    callsites.values.toList.foreach(rewriteFinalTraitMethodInvocation)
+    callsites.valuesIterator.flatMap(_.valuesIterator).toList.foreach(rewriteFinalTraitMethodInvocation)
   }
 
   /**
@@ -202,7 +202,7 @@ class Inliner[BT <: BTypes](val btypes: BT) {
         callsite.callsiteMethod.instructions.insert(callsite.callsiteInstruction, newCallsiteInstruction)
         callsite.callsiteMethod.instructions.remove(callsite.callsiteInstruction)
 
-        callGraph.callsites.remove(callsite.callsiteInstruction)
+        callGraph.removeCallsite(callsite.callsiteInstruction, callsite.callsiteMethod)
         val staticCallsite = Callsite(
           callsiteInstruction = newCallsiteInstruction,
           callsiteMethod      = callsite.callsiteMethod,
@@ -220,13 +220,13 @@ class Inliner[BT <: BTypes](val btypes: BT) {
           receiverKnownNotNull = callsite.receiverKnownNotNull,
           callsitePosition = callsite.callsitePosition
         )
-        callGraph.callsites(newCallsiteInstruction) = staticCallsite
+        callGraph.addCallsite(staticCallsite)
       }
 
       for (warning <- res.left) {
         val Right(callee) = callsite.callee
         val newCallee = callee.copy(calleeInfoWarning = Some(RewriteTraitCallToStaticImplMethodFailed(calleeDeclarationClass.internalName, callee.callee.name, callee.callee.desc, warning)))
-        callGraph.callsites(callsite.callsiteInstruction) = callsite.copy(callee = Right(newCallee))
+        callGraph.addCallsite(callsite.copy(callee = Right(newCallee)))
       }
     }
   }
@@ -435,38 +435,30 @@ class Inliner[BT <: BTypes](val btypes: BT) {
       callsiteMethod.tryCatchBlocks.addAll(cloneTryCatchBlockNodes(callee, labelsMap).asJava)
 
       // Add all invocation instructions and closure instantiations that were inlined to the call graph
-      callee.instructions.iterator().asScala foreach {
-        case originalCallsiteIns: MethodInsnNode =>
-          callGraph.callsites.get(originalCallsiteIns) match {
-            case Some(originalCallsite) =>
-              val newCallsiteIns = instructionMap(originalCallsiteIns).asInstanceOf[MethodInsnNode]
-              callGraph.callsites(newCallsiteIns) = Callsite(
-                callsiteInstruction = newCallsiteIns,
-                callsiteMethod = callsiteMethod,
-                callsiteClass = callsiteClass,
-                callee = originalCallsite.callee,
-                argInfos = Nil, // TODO: re-compute argInfos for new destination (once we actually compute them)
-                callsiteStackHeight = callsiteStackHeight + originalCallsite.callsiteStackHeight,
-                receiverKnownNotNull = originalCallsite.receiverKnownNotNull,
-                callsitePosition = originalCallsite.callsitePosition
-              )
-
-            case None =>
-          }
-
-        case indy: InvokeDynamicInsnNode =>
-          callGraph.closureInstantiations.get(indy) match {
-            case Some(closureInit) =>
-              val newIndy = instructionMap(indy).asInstanceOf[InvokeDynamicInsnNode]
-              callGraph.closureInstantiations(newIndy) = ClosureInstantiation(closureInit.lambdaMetaFactoryCall.copy(indy = newIndy), callsiteMethod, callsiteClass)
-
-            case None =>
-          }
+      callGraph.callsites(callee).valuesIterator foreach { originalCallsite =>
+        val newCallsiteIns = instructionMap(originalCallsite.callsiteInstruction).asInstanceOf[MethodInsnNode]
+        callGraph.addCallsite(Callsite(
+          callsiteInstruction = newCallsiteIns,
+          callsiteMethod = callsiteMethod,
+          callsiteClass = callsiteClass,
+          callee = originalCallsite.callee,
+          argInfos = Nil, // TODO: re-compute argInfos for new destination (once we actually compute them)
+          callsiteStackHeight = callsiteStackHeight + originalCallsite.callsiteStackHeight,
+          receiverKnownNotNull = originalCallsite.receiverKnownNotNull,
+          callsitePosition = originalCallsite.callsitePosition
+          )
+        )
+      }
 
-        case _ =>
+      callGraph.closureInstantiations(callee).valuesIterator foreach { originalClosureInit =>
+        val newIndy = instructionMap(originalClosureInit.lambdaMetaFactoryCall.indy).asInstanceOf[InvokeDynamicInsnNode]
+        callGraph.addClosureInstantiation(
+          ClosureInstantiation(originalClosureInit.lambdaMetaFactoryCall.copy(indy = newIndy), callsiteMethod, callsiteClass)
+        )
       }
+
       // Remove the elided invocation from the call graph
-      callGraph.callsites.remove(callsiteInstruction)
+      callGraph.removeCallsite(callsiteInstruction, callsiteMethod)
 
       // Inlining a method body can render some code unreachable, see example above (in runInliner).
       unreachableCodeEliminated -= callsiteMethod