Rewrite closure invocations to the lambda body method

When an indylambda closure is allocated and invoked within the same
method, rewrite the invocation to the implementation method.

This works for any indylambda / SAM type, not only Scala functions.
However, the Scala compiler (under -Xexperimental) currently desugars
function literals for non-FunctionN types to an anonymous class during
typer.

No testing yet, waiting for FunctionN to become SAMs first.

The feature requires scala-java8-compat to be on the classpath and a
number of compiler flags:
-Ydelambdafy:method -Ybackend:GenBCode -Yopt:closure-elimination -target:jvm-1.8

➜  scala git:(opt/closureInlining) ant -Dscala-java8-compat.package=1 -Dlocker.skip=1
➜  scala git:(opt/closureInlining) cd sandbox

➜  sandbox git:(opt/closureInlining) cat Fun.java
public interface Fun<T> {
  T apply(T x);
}
➜  sandbox git:(opt/closureInlining) javac Fun.java

➜  sandbox git:(opt/closureInlining) cat Test.scala
class C {
  val z = "too"
  def f = {
    val kap = "me! me!"
    val f: Tuple2[String, String] => String = (o => z + kap + o.toString)
    f(("a", "b"))
  }

  def g = {
    val f: Int => String = x => x.toString
    f(10)
  }

  def h = {
    val f: Fun[Int] = x => x + 100 // Java SAM, requires -Xexperimental, will create an anonymous class in typer
    f(10)
  }

  def i = {
    val l = 10l
    val f: (Long, String) => String = (x, s) => s + l + z + x
    f(20l, "n")
  }

  def j = {
    val f: Int => Int = x => x + 101 // specialized
    f(33)
  }
}
➜  sandbox git:(opt/closureInlining) ../build/quick/bin/scalac -target:jvm-1.8 -Yopt:closure-elimination -Ydelambdafy:method -Ybackend:GenBCode -Xexperimental -cp ../build/quick/scala-java8-compat:. Test.scala

➜  sandbox git:(opt/closureInlining) asm -a C.class
➜  sandbox git:(opt/closureInlining) cat C.asm

[...]

  public g()Ljava/lang/String;
   L0
    INVOKEDYNAMIC apply()Lscala/compat/java8/JFunction1; [
      // handle kind 0x6 : INVOKESTATIC
      java/lang/invoke/LambdaMetafactory.altMetafactory(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;[Ljava/lang/Object;)Ljava/lang/invoke/CallSite;
      // arguments:
      (Ljava/lang/Object;)Ljava/lang/Object;,
      // handle kind 0x6 : INVOKESTATIC
      C.C$$$anonfun$2$adapted(Ljava/lang/Object;)Ljava/lang/String;,
      (Ljava/lang/Object;)Ljava/lang/String;,
      3,
      1,
      Lscala/Serializable;.class,
      0
    ]
    CHECKCAST scala/Function1
   L1
    ASTORE 1
   L2
    ALOAD 1
    BIPUSH 10
    INVOKESTATIC scala/runtime/BoxesRunTime.boxToInteger (I)Ljava/lang/Integer;
    ASTORE 2
    POP
    ALOAD 2
    INVOKESTATIC C.C$$$anonfun$2$adapted (Ljava/lang/Object;)Ljava/lang/String;
    CHECKCAST java/lang/String
   L3
    ARETURN

[...]

Author: lrytz

src/compiler/scala/tools/nsc/backend/jvm/BCodeSkelBuilder.scala

@@ -140,7 +140,7 @@ abstract class BCodeSkelBuilder extends BCodeHelpers {
       if (AsmUtils.traceClassEnabled && cnode.name.contains(AsmUtils.traceClassPattern))
         AsmUtils.traceClass(cnode)
 
-      if (settings.YoptInlinerEnabled) {
+      if (settings.YoptAddToBytecodeRepository) {
         // The inliner needs to find all classes in the code repo, also those being compiled
         byteCodeRepository.add(cnode, ByteCodeRepository.CompilationUnit)
       }

src/compiler/scala/tools/nsc/backend/jvm/BTypes.scala

@@ -44,6 +44,8 @@ abstract class BTypes {
 
   val inliner: Inliner[this.type]
 
+  val closureOptimizer: ClosureOptimizer[this.type]
+
   val callGraph: CallGraph[this.type]
 
   val backendReporting: BackendReporting

src/compiler/scala/tools/nsc/backend/jvm/BTypesFromSymbols.scala

@@ -7,7 +7,7 @@ package scala.tools.nsc
 package backend.jvm
 
 import scala.tools.asm
-import scala.tools.nsc.backend.jvm.opt.{LocalOpt, CallGraph, Inliner, ByteCodeRepository}
+import scala.tools.nsc.backend.jvm.opt._
 import scala.tools.nsc.backend.jvm.BTypes.{InlineInfo, MethodInlineInfo, InternalName}
 import BackendReporting._
 import scala.tools.nsc.settings.ScalaSettings
@@ -42,6 +42,8 @@ class BTypesFromSymbols[G <: Global](val global: G) extends BTypes {
 
   val inliner: Inliner[this.type] = new Inliner(this)
 
+  val closureOptimizer: ClosureOptimizer[this.type] = new ClosureOptimizer(this)
+
   val callGraph: CallGraph[this.type] = new CallGraph(this)
 
   val backendReporting: BackendReporting = new BackendReportingImpl(global)

src/compiler/scala/tools/nsc/backend/jvm/BackendReporting.scala

@@ -246,6 +246,28 @@ object BackendReporting {
   case class ResultingMethodTooLarge(calleeDeclarationClass: InternalName, name: String, descriptor: String,
                                      callsiteClass: InternalName, callsiteName: String, callsiteDesc: String) extends CannotInlineWarning
 
+  /**
+   * Used in `rewriteClosureApplyInvocations` when a closure apply callsite cannot be rewritten
+   * to the closure body method.
+   */
+  trait RewriteClosureApplyToClosureBodyFailed extends OptimizerWarning {
+    def pos: Position
+
+    override def emitWarning(settings: ScalaSettings): Boolean = this match {
+      case RewriteClosureAccessCheckFailed(_, cause) => cause.emitWarning(settings)
+      case RewriteClosureIllegalAccess(_, _)         => settings.YoptWarningEmitAtInlineFailed
+    }
+
+    override def toString: String = this match {
+      case RewriteClosureAccessCheckFailed(_, cause) =>
+        s"Failed to rewrite the closure invocation to its implementation method:\n" + cause
+      case RewriteClosureIllegalAccess(_, callsiteClass) =>
+        s"The closure body invocation cannot be rewritten because the target method is not accessible in class $callsiteClass."
+    }
+  }
+  case class RewriteClosureAccessCheckFailed(pos: Position, cause: OptimizerWarning) extends RewriteClosureApplyToClosureBodyFailed
+  case class RewriteClosureIllegalAccess(pos: Position, callsiteClass: InternalName) extends RewriteClosureApplyToClosureBodyFailed
+
   /**
    * Used in the InlineInfo of a ClassBType, when some issue occurred obtaining the inline information.
    */

src/compiler/scala/tools/nsc/backend/jvm/GenBCode.scala

@@ -216,20 +216,25 @@ abstract class GenBCode extends BCodeSyncAndTry {
     class Worker2 {
       def runGlobalOptimizations(): Unit = {
         import scala.collection.convert.decorateAsScala._
-        q2.asScala foreach {
-          case Item2(_, _, plain, _, _) =>
-            // skip mirror / bean: wd don't inline into tem, and they are not used in the plain class
-            if (plain != null) callGraph.addClass(plain)
+        if (settings.YoptBuildCallGraph) {
+          q2.asScala foreach {
+            case Item2(_, _, plain, _, _) =>
+              // skip mirror / bean: wd don't inline into tem, and they are not used in the plain class
+              if (plain != null) callGraph.addClass(plain)
+          }
         }
-        bTypes.inliner.runInliner()
+        if (settings.YoptInlinerEnabled)
+          bTypes.inliner.runInliner()
+        if (settings.YoptClosureElimination)
+          closureOptimizer.rewriteClosureApplyInvocations()
       }
 
       def localOptimizations(classNode: ClassNode): Unit = {
         BackendStats.timed(BackendStats.methodOptTimer)(localOpt.methodOptimizations(classNode))
       }
 
       def run() {
-        if (settings.YoptInlinerEnabled) runGlobalOptimizations()
+        runGlobalOptimizations()
 
         while (true) {
           val item = q2.poll

src/compiler/scala/tools/nsc/backend/jvm/opt/ByteCodeRepository.scala

@@ -102,7 +102,7 @@ class ByteCodeRepository(val classPath: ClassFileLookup[AbstractFile], val isJav
   }
 
   /**
-   * The method node for a method matching `name` and `descriptor`, accessed in class `classInternalName`.
+   * The method node for a method matching `name` and `descriptor`, accessed in class `ownerInternalNameOrArrayDescriptor`.
    * The declaration of the method may be in one of the parents.
    *
    * @return The [[MethodNode]] of the requested method and the [[InternalName]] of its declaring

src/compiler/scala/tools/nsc/backend/jvm/opt/CallGraph.scala

@@ -11,6 +11,7 @@ import scala.reflect.internal.util.{NoPosition, Position}
 import scala.tools.asm.tree.analysis.{Value, Analyzer, BasicInterpreter}
 import scala.tools.asm.{Opcodes, Type}
 import scala.tools.asm.tree._
+import scala.collection.concurrent
 import scala.collection.convert.decorateAsScala._
 import scala.tools.nsc.backend.jvm.BTypes.InternalName
 import scala.tools.nsc.backend.jvm.BackendReporting._
@@ -21,14 +22,25 @@ import BytecodeUtils._
 class CallGraph[BT <: BTypes](val btypes: BT) {
   import btypes._
 
-  val callsites: collection.concurrent.Map[MethodInsnNode, Callsite] = recordPerRunCache(collection.concurrent.TrieMap.empty[MethodInsnNode, Callsite])
+  val callsites: concurrent.Map[MethodInsnNode, Callsite] = recordPerRunCache(concurrent.TrieMap.empty)
+
+  val closureInstantiations: concurrent.Map[InvokeDynamicInsnNode, (MethodNode, ClassBType)] = recordPerRunCache(concurrent.TrieMap.empty)
 
   def addClass(classNode: ClassNode): Unit = {
-    for (m <- classNode.methods.asScala; callsite <- analyzeCallsites(m, classBTypeFromClassNode(classNode)))
-      callsites(callsite.callsiteInstruction) = callsite
+    val classType = classBTypeFromClassNode(classNode)
+    for {
+      m <- classNode.methods.asScala
+      (calls, closureInits) = analyzeCallsites(m, classType)
+    } {
+      calls foreach (callsite => callsites(callsite.callsiteInstruction) = callsite)
+      closureInits foreach (indy => closureInstantiations(indy) = (m, classType))
+    }
   }
 
-  def analyzeCallsites(methodNode: MethodNode, definingClass: ClassBType): List[Callsite] = {
+  /**
+   * 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[InvokeDynamicInsnNode]) = {
 
     case class CallsiteInfo(safeToInline: Boolean, safeToRewrite: Boolean,
                             annotatedInline: Boolean, annotatedNoInline: Boolean,
@@ -116,7 +128,10 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
       case _ => false
     }
 
-    methodNode.instructions.iterator.asScala.collect({
+    val callsites = new collection.mutable.ListBuffer[Callsite]
+    val closureInstantiations = new collection.mutable.ListBuffer[InvokeDynamicInsnNode]
+
+    methodNode.instructions.iterator.asScala foreach {
       case call: MethodInsnNode =>
         val callee: Either[OptimizerWarning, Callee] = for {
           (method, declarationClass)     <- byteCodeRepository.methodNode(call.owner, call.name, call.desc): Either[OptimizerWarning, (MethodNode, InternalName)]
@@ -147,7 +162,7 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
           receiverNotNullByAnalysis(call, numArgs)
         }
 
-        Callsite(
+        callsites += Callsite(
           callsiteInstruction = call,
           callsiteMethod = methodNode,
           callsiteClass = definingClass,
@@ -157,7 +172,14 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
           receiverKnownNotNull = receiverNotNull,
           callsitePosition = callsitePositions.getOrElse(call, NoPosition)
         )
-    }).toList
+
+      case indy: InvokeDynamicInsnNode =>
+        if (closureOptimizer.isClosureInstantiation(indy)) closureInstantiations += indy
+
+      case _ =>
+    }
+
+    (callsites.toList, closureInstantiations.toList)
   }
 
   /**
@@ -201,7 +223,7 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
    * @param calleeDeclarationClass The class in which the callee is declared
    * @param safeToInline           True if the callee can be safely inlined: it cannot be overridden,
    *                               and the inliner settings (project / global) allow inlining it.
-   * @param safeToRewrite          True if the callee the interface method of a concrete trait method
+   * @param safeToRewrite          True if the callee is the interface method of a concrete trait method
    *                               that can be safely re-written to the static implementation method.
    * @param annotatedInline        True if the callee is annotated @inline
    * @param annotatedNoInline      True if the callee is annotated @noinline