Refactor the ClosureOptimizer, run ProdCons only once per method

Refactor and clean up the ClosureOptimzier. The goal of this change
is to run the ProdCons analyzer only once per method, instead of
once per closure instantiation.

Bootstrapping scala with `-Yopt-inline-heuristics:everything` revealed
that ProdCons can take a very long time on large methods, for example

```
[quick.compiler] scala/tools/nsc/backend/jvm/BCodeBodyBuilder$PlainBodyBuilder#genArithmeticOp - analysis: 1 spans, 17755ms
[quick.compiler] scala/tools/nsc/typechecker/SuperAccessors$SuperAccTransformer#transform - analysis: 1 spans, 28024ms
[quick.compiler] scala/tools/nsc/backend/jvm/BCodeBodyBuilder$PlainBodyBuilder#genInvokeDynamicLambda - analysis: 1 spans, 22100ms
```

With this change and enough time and space (-Xmx6000m), bootstrapping
scala succeeds in this test mode.

1 files changed, 210 insertions, 157 deletions

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 86536ff0d2..58f0813bd8 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/ClosureOptimizer.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/ClosureOptimizer.scala
@@ -8,9 +8,9 @@ package backend.jvm
 package opt
 
 import scala.annotation.switch
-import scala.collection.mutable
+import scala.collection.immutable
 import scala.reflect.internal.util.NoPosition
-import scala.tools.asm.{Handle, Type, Opcodes}
+import scala.tools.asm.{Type, Opcodes}
 import scala.tools.asm.tree._
 import scala.tools.nsc.backend.jvm.BTypes.InternalName
 import scala.tools.nsc.backend.jvm.analysis.ProdConsAnalyzer
@@ -24,6 +24,35 @@ class ClosureOptimizer[BT <: BTypes](val btypes: BT) {
   import btypes._
   import callGraph._
 
+  /**
+   * If a closure is allocated and invoked within the same method, re-write the invocation to the
+   * closure body method.
+   *
+   * Note that the closure body method (generated by delambdafy:method) takes additional parameters
+   * for the values captured by the closure. The bytecode is transformed from
+   *
+   *   [generate captured values]
+   *   [closure init, capturing values]
+   *   [...]
+   *   [load closure object]
+   *   [generate closure invocation arguments]
+   *   [invoke closure.apply]
+   *
+   * to
+   *
+   *   [generate captured values]
+   *   [store captured values into new locals]
+   *   [load the captured values from locals]    // a future optimization will eliminate the closure
+   *   [closure init, capturing values]          // instantiation if the closure object becomes unused
+   *   [...]
+   *   [load closure object]
+   *   [generate closure invocation arguments]
+   *   [store argument values into new locals]
+   *   [drop the closure object]
+   *   [load captured values from locals]
+   *   [load argument values from locals]
+   *   [invoke the closure body method]
+   */
   def rewriteClosureApplyInvocations(): Unit = {
     implicit object closureInitOrdering extends Ordering[ClosureInstantiation] {
       override def compare(x: ClosureInstantiation, y: ClosureInstantiation): Int = {
@@ -41,16 +70,111 @@ class ClosureOptimizer[BT <: BTypes](val btypes: BT) {
       }
     }
 
-    val sorted = mutable.TreeSet.empty[ClosureInstantiation]
-    sorted ++= closureInstantiations.values
+    // Group 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 (closureInst <- sorted) {
-      val warnings = rewriteClosureApplyInvocations(closureInst)
-      warnings.foreach(w => backendReporting.inlinerWarning(w.pos, w.toString))
+    // 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: Map[ClosureInstantiation, List[Either[RewriteClosureApplyToClosureBodyFailed, (MethodInsnNode, Int)]]] = {
+      closureInstantiationsByMethod 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.map(init => (init, closureCallsites(init, prodCons)))
+      }
+    }
+
+    // Rewrite all closure callsites (or issue inliner warnings for those that cannot be rewritten)
+    for ((closureInit, callsites) <- callsitesToRewrite) {
+      // Local variables that hold the captured values and the closure invocation arguments.
+      // They are lazy vals to ensure that locals for captured values are only allocated if there's
+      // actually a callsite to rewrite (an not only warnings to be issued).
+      lazy val (localsForCapturedValues, argumentLocalsList) = localsForClosureRewrite(closureInit)
+      for (callsite <- callsites) callsite match {
+        case Left(warning) =>
+          backendReporting.inlinerWarning(warning.pos, warning.toString)
+
+        case Right((invocation, stackHeight)) =>
+          rewriteClosureApplyInvocation(closureInit, invocation, stackHeight, localsForCapturedValues, argumentLocalsList)
+      }
     }
   }
 
-  def isSamInvocation(invocation: MethodInsnNode, indy: InvokeDynamicInsnNode, prodCons: => ProdConsAnalyzer): Boolean = {
+  /**
+   * Insert instructions to store the values captured by a closure instantiation into local variables,
+   * and load the values back to the stack.
+   *
+   * Returns the list of locals holding those captured values, and a list of locals that should be
+   * used at the closure invocation callsite to store the arguments passed to the closure invocation.
+   */
+  private def localsForClosureRewrite(closureInit: ClosureInstantiation): (LocalsList, LocalsList) = {
+    val ownerMethod = closureInit.ownerMethod
+    val captureLocals = storeCaptures(closureInit)
+
+    // allocate locals for storing the arguments of the closure apply callsites.
+    // if there are multiple callsites, the same locals are re-used.
+    val argTypes = closureInit.lambdaMetaFactoryCall.samMethodType.getArgumentTypes
+    val firstArgLocal = ownerMethod.maxLocals
+
+    // The comment in the unapply method of `LambdaMetaFactoryCall` explains why we have to introduce
+    // casts for arguments that have different types in samMethodType and instantiatedMethodType.
+    val castLoadTypes = {
+      val instantiatedMethodType = closureInit.lambdaMetaFactoryCall.instantiatedMethodType
+      (argTypes, instantiatedMethodType.getArgumentTypes).zipped map {
+        case (samArgType, instantiatedArgType) if samArgType != instantiatedArgType =>
+          // the LambdaMetaFactoryCall extractor ensures that the two types are reference types,
+          // so we don't end up casting primitive values.
+          Some(instantiatedArgType)
+        case _ =>
+          None
+      }
+    }
+    val argLocals = LocalsList.fromTypes(firstArgLocal, argTypes, castLoadTypes)
+    ownerMethod.maxLocals = firstArgLocal + argLocals.size
+
+    (captureLocals, argLocals)
+  }
+
+  /**
+   * Find all callsites of a closure within the method where the closure is allocated.
+   */
+  private def closureCallsites(closureInit: ClosureInstantiation, prodCons: => ProdConsAnalyzer): List[Either[RewriteClosureApplyToClosureBodyFailed, (MethodInsnNode, Int)]] = {
+    val ownerMethod = closureInit.ownerMethod
+    val ownerClass = closureInit.ownerClass
+    val lambdaBodyHandle = closureInit.lambdaMetaFactoryCall.implMethod
+
+    ownerMethod.instructions.iterator.asScala.collect({
+      case invocation: MethodInsnNode if isSamInvocation(invocation, closureInit, prodCons) =>
+        // TODO: This is maybe over-cautious.
+        // We are checking if the closure body method is accessible at the closure callsite.
+        // If the closure allocation has access to the body method, then the callsite (in the same
+        // method as the alloction) should have access too.
+        val bodyAccessible: Either[OptimizerWarning, Boolean] = for {
+          (bodyMethodNode, declClass) <- byteCodeRepository.methodNode(lambdaBodyHandle.getOwner, lambdaBodyHandle.getName, lambdaBodyHandle.getDesc): Either[OptimizerWarning, (MethodNode, InternalName)]
+          isAccessible                <- inliner.memberIsAccessible(bodyMethodNode.access, classBTypeFromParsedClassfile(declClass), classBTypeFromParsedClassfile(lambdaBodyHandle.getOwner), ownerClass)
+        } yield {
+          isAccessible
+        }
+
+        def pos = callGraph.callsites.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))
+          case _            => Right(prodCons.frameAt(invocation).getStackSize)
+        }
+
+        stackSize.right.map((invocation, _))
+    }).toList
+  }
+
+  private def isSamInvocation(invocation: MethodInsnNode, closureInit: ClosureInstantiation, prodCons: => ProdConsAnalyzer): Boolean = {
+    val indy = closureInit.lambdaMetaFactoryCall.indy
     if (invocation.getOpcode == INVOKESTATIC) false
     else {
       def closureIsReceiver = {
@@ -64,16 +188,90 @@ class ClosureOptimizer[BT <: BTypes](val btypes: BT) {
       }
 
       invocation.name == indy.name && {
-        val indySamMethodDesc = indy.bsmArgs(0).asInstanceOf[Type].getDescriptor // safe, checked in isClosureInstantiation
+        val indySamMethodDesc = closureInit.lambdaMetaFactoryCall.samMethodType.getDescriptor
         indySamMethodDesc == invocation.desc
       } &&
-      closureIsReceiver // most expensive check last
+        closureIsReceiver // most expensive check last
     }
   }
 
+  private def rewriteClosureApplyInvocation(closureInit: ClosureInstantiation, invocation: MethodInsnNode, stackHeight: Int, localsForCapturedValues: LocalsList, argumentLocalsList: LocalsList): Unit = {
+    val ownerMethod = closureInit.ownerMethod
+    val lambdaBodyHandle = closureInit.lambdaMetaFactoryCall.implMethod
+
+    // store arguments
+    insertStoreOps(invocation, ownerMethod, argumentLocalsList)
+
+    // drop the closure from the stack
+    ownerMethod.instructions.insertBefore(invocation, new InsnNode(POP))
+
+    // load captured values and arguments
+    insertLoadOps(invocation, ownerMethod, localsForCapturedValues)
+    insertLoadOps(invocation, ownerMethod, argumentLocalsList)
+
+    // update maxStack
+    val capturesStackSize = localsForCapturedValues.size
+    val invocationStackHeight = stackHeight + capturesStackSize - 1 // -1 because the closure is gone
+    if (invocationStackHeight > ownerMethod.maxStack)
+      ownerMethod.maxStack = invocationStackHeight
+
+    // replace the callsite with a new call to the body method
+    val bodyOpcode = (lambdaBodyHandle.getTag: @switch) match {
+      case H_INVOKEVIRTUAL    => INVOKEVIRTUAL
+      case H_INVOKESTATIC     => INVOKESTATIC
+      case H_INVOKESPECIAL    => INVOKESPECIAL
+      case H_INVOKEINTERFACE  => INVOKEINTERFACE
+      case H_NEWINVOKESPECIAL =>
+        val insns = ownerMethod.instructions
+        insns.insertBefore(invocation, new TypeInsnNode(NEW, lambdaBodyHandle.getOwner))
+        insns.insertBefore(invocation, new InsnNode(DUP))
+        INVOKESPECIAL
+    }
+    val isInterface = bodyOpcode == INVOKEINTERFACE
+    val bodyInvocation = new MethodInsnNode(bodyOpcode, lambdaBodyHandle.getOwner, lambdaBodyHandle.getName, lambdaBodyHandle.getDesc, isInterface)
+    ownerMethod.instructions.insertBefore(invocation, bodyInvocation)
+
+    val returnType = Type.getReturnType(lambdaBodyHandle.getDesc)
+    fixLoadedNothingOrNullValue(returnType, bodyInvocation, ownerMethod, btypes) // see comment of that method
+
+    ownerMethod.instructions.remove(invocation)
+
+    // update the call graph
+    val originalCallsite = callGraph.callsites.remove(invocation)
+
+    // the method node is needed for building the call graph entry
+    val bodyMethod = byteCodeRepository.methodNode(lambdaBodyHandle.getOwner, lambdaBodyHandle.getName, lambdaBodyHandle.getDesc)
+    def bodyMethodIsBeingCompiled = byteCodeRepository.classNodeAndSource(lambdaBodyHandle.getOwner).map(_._2 == CompilationUnit).getOrElse(false)
+    val bodyMethodCallsite = Callsite(
+      callsiteInstruction = bodyInvocation,
+      callsiteMethod = ownerMethod,
+      callsiteClass = closureInit.ownerClass,
+      callee = bodyMethod.map({
+        case (bodyMethodNode, bodyMethodDeclClass) => Callee(
+          callee = bodyMethodNode,
+          calleeDeclarationClass = classBTypeFromParsedClassfile(bodyMethodDeclClass),
+          safeToInline = compilerSettings.YoptInlineGlobal || bodyMethodIsBeingCompiled,
+          safeToRewrite = false, // the lambda body method is not a trait interface method
+          annotatedInline = false,
+          annotatedNoInline = false,
+          calleeInfoWarning = None)
+      }),
+      argInfos = Nil,
+      callsiteStackHeight = invocationStackHeight,
+      receiverKnownNotNull = true, // see below (*)
+      callsitePosition = originalCallsite.map(_.callsitePosition).getOrElse(NoPosition)
+    )
+    // (*) The documentation in class LambdaMetafactory says:
+    //     "if implMethod corresponds to an instance method, the first capture argument
+    //     (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
+  }
+
   /**
-   * Stores the values captured by a closure creation into fresh local variables.
-   * Returns the list of locals holding the captured values.
+   * Stores the values captured by a closure creation into fresh local variables, and loads the
+   * values back onto the stack. Returns the list of locals holding the captured values.
    */
   private def storeCaptures(closureInit: ClosureInstantiation): LocalsList = {
     val indy = closureInit.lambdaMetaFactoryCall.indy
@@ -129,151 +327,6 @@ class ClosureOptimizer[BT <: BTypes](val btypes: BT) {
     }
   }
 
-  def rewriteClosureApplyInvocations(closureInit: ClosureInstantiation): List[RewriteClosureApplyToClosureBodyFailed] = {
-    val lambdaBodyHandle = closureInit.lambdaMetaFactoryCall.implMethod
-    val ownerMethod = closureInit.ownerMethod
-    val ownerClass = closureInit.ownerClass
-
-    // Kept as a lazy val to make sure the analysis is only computed if it's actually needed.
-    // ProdCons is used to identify closure body invocations (see isSamInvocation), but only if the
-    // callsite has the right name and signature. If the method has no invcation instruction with
-    // the right name and signature, the analysis is not executed.
-    lazy val prodCons = new ProdConsAnalyzer(ownerMethod, ownerClass.internalName)
-
-    // First collect all callsites without modifying the instructions list yet.
-    // Once we start modifying the instruction list, prodCons becomes unusable.
-
-    // A list of callsites and stack heights. If the invocation cannot be rewritten, a warning
-    // message is stored in the stack height value.
-    val invocationsToRewrite: List[(MethodInsnNode, Either[RewriteClosureApplyToClosureBodyFailed, Int])] = ownerMethod.instructions.iterator.asScala.collect({
-      case invocation: MethodInsnNode if isSamInvocation(invocation, closureInit.lambdaMetaFactoryCall.indy, prodCons) =>
-        val bodyAccessible: Either[OptimizerWarning, Boolean] = for {
-          (bodyMethodNode, declClass) <- byteCodeRepository.methodNode(lambdaBodyHandle.getOwner, lambdaBodyHandle.getName, lambdaBodyHandle.getDesc): Either[OptimizerWarning, (MethodNode, InternalName)]
-          isAccessible                <- inliner.memberIsAccessible(bodyMethodNode.access, classBTypeFromParsedClassfile(declClass), classBTypeFromParsedClassfile(lambdaBodyHandle.getOwner), ownerClass)
-        } yield {
-          isAccessible
-        }
-
-        def pos = callGraph.callsites.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))
-          case _            => Right(prodCons.frameAt(invocation).getStackSize)
-        }
-
-        (invocation, stackSize)
-    }).toList
-
-    if (invocationsToRewrite.isEmpty) Nil
-    else {
-      // lazy val to make sure locals for captures and arguments are only allocated if there's
-      // effectively a callsite to rewrite.
-      lazy val (localsForCapturedValues, argumentLocalsList) = {
-        val captureLocals = storeCaptures(closureInit)
-
-        // allocate locals for storing the arguments of the closure apply callsites.
-        // if there are multiple callsites, the same locals are re-used.
-        val argTypes = closureInit.lambdaMetaFactoryCall.samMethodType.getArgumentTypes
-        val firstArgLocal = ownerMethod.maxLocals
-
-        // The comment in `isClosureInstantiation` explains why we have to introduce casts for
-        // arguments that have different types in samMethodType and instantiatedMethodType.
-        val castLoadTypes = {
-          val instantiatedMethodType = closureInit.lambdaMetaFactoryCall.instantiatedMethodType
-          (argTypes, instantiatedMethodType.getArgumentTypes).zipped map {
-            case (samArgType, instantiatedArgType) if samArgType != instantiatedArgType =>
-              // isClosureInstantiation ensures that the two types are reference types, so we don't
-              // end up casting primitive values.
-              Some(instantiatedArgType)
-            case _ =>
-              None
-          }
-        }
-        val argLocals = LocalsList.fromTypes(firstArgLocal, argTypes, castLoadTypes)
-        ownerMethod.maxLocals = firstArgLocal + argLocals.size
-
-        (captureLocals, argLocals)
-      }
-
-      val warnings = invocationsToRewrite flatMap {
-        case (invocation, Left(warning)) => Some(warning)
-
-        case (invocation, Right(stackHeight)) =>
-          // store arguments
-          insertStoreOps(invocation, ownerMethod, argumentLocalsList)
-
-          // drop the closure from the stack
-          ownerMethod.instructions.insertBefore(invocation, new InsnNode(POP))
-
-          // load captured values and arguments
-          insertLoadOps(invocation, ownerMethod, localsForCapturedValues)
-          insertLoadOps(invocation, ownerMethod, argumentLocalsList)
-
-          // update maxStack
-          val capturesStackSize = localsForCapturedValues.size
-          val invocationStackHeight = stackHeight + capturesStackSize - 1 // -1 because the closure is gone
-          if (invocationStackHeight > ownerMethod.maxStack)
-            ownerMethod.maxStack = invocationStackHeight
-
-          // replace the callsite with a new call to the body method
-          val bodyOpcode = (lambdaBodyHandle.getTag: @switch) match {
-            case H_INVOKEVIRTUAL    => INVOKEVIRTUAL
-            case H_INVOKESTATIC     => INVOKESTATIC
-            case H_INVOKESPECIAL    => INVOKESPECIAL
-            case H_INVOKEINTERFACE  => INVOKEINTERFACE
-            case H_NEWINVOKESPECIAL =>
-              val insns = ownerMethod.instructions
-              insns.insertBefore(invocation, new TypeInsnNode(NEW, lambdaBodyHandle.getOwner))
-              insns.insertBefore(invocation, new InsnNode(DUP))
-              INVOKESPECIAL
-          }
-          val isInterface = bodyOpcode == INVOKEINTERFACE
-          val bodyInvocation = new MethodInsnNode(bodyOpcode, lambdaBodyHandle.getOwner, lambdaBodyHandle.getName, lambdaBodyHandle.getDesc, isInterface)
-          ownerMethod.instructions.insertBefore(invocation, bodyInvocation)
-
-          val returnType = Type.getReturnType(lambdaBodyHandle.getDesc)
-          fixLoadedNothingOrNullValue(returnType, bodyInvocation, ownerMethod, btypes) // see comment of that method
-
-          ownerMethod.instructions.remove(invocation)
-
-          // update the call graph
-          val originalCallsite = callGraph.callsites.remove(invocation)
-
-          // the method node is needed for building the call graph entry
-          val bodyMethod = byteCodeRepository.methodNode(lambdaBodyHandle.getOwner, lambdaBodyHandle.getName, lambdaBodyHandle.getDesc)
-          def bodyMethodIsBeingCompiled = byteCodeRepository.classNodeAndSource(lambdaBodyHandle.getOwner).map(_._2 == CompilationUnit).getOrElse(false)
-          val bodyMethodCallsite = Callsite(
-            callsiteInstruction = bodyInvocation,
-            callsiteMethod = ownerMethod,
-            callsiteClass = ownerClass,
-            callee = bodyMethod.map({
-              case (bodyMethodNode, bodyMethodDeclClass) => Callee(
-                callee = bodyMethodNode,
-                calleeDeclarationClass = classBTypeFromParsedClassfile(bodyMethodDeclClass),
-                safeToInline = compilerSettings.YoptInlineGlobal || bodyMethodIsBeingCompiled,
-                safeToRewrite = false, // the lambda body method is not a trait interface method
-                annotatedInline = false,
-                annotatedNoInline = false,
-                calleeInfoWarning = None)
-            }),
-            argInfos = Nil,
-            callsiteStackHeight = invocationStackHeight,
-            receiverKnownNotNull = true, // see below (*)
-            callsitePosition = originalCallsite.map(_.callsitePosition).getOrElse(NoPosition)
-          )
-          // (*) The documentation in class LambdaMetafactory says:
-          //     "if implMethod corresponds to an instance method, the first capture argument
-          //     (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
-          None
-      }
-
-      warnings.toList
-    }
-  }
-
   /**
    * A list of local variables. Each local stores information about its type, see class [[Local]].
    */