Simplify post inlining requests

Clean up inliner test

3 files changed, 34 insertions, 22 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 a8f1e43071..801296908f 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/CallGraph.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/CallGraph.scala
@@ -67,6 +67,8 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
     callsites(callsite.callsiteMethod) = methodCallsites + (callsite.callsiteInstruction -> callsite)
   }
 
+  def containsCallsite(callsite: Callsite): Boolean = callsites(callsite.callsiteMethod) contains callsite.callsiteInstruction
+
   def removeClosureInstantiation(indy: InvokeDynamicInsnNode, methodNode: MethodNode): Option[ClosureInstantiation] = {
     val methodClosureInits = closureInstantiations(methodNode)
     val newClosureInits = methodClosureInits - indy
@@ -130,8 +132,8 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
           val callee: Either[OptimizerWarning, Callee] = for {
             (method, declarationClass)     <- byteCodeRepository.methodNode(call.owner, call.name, call.desc): Either[OptimizerWarning, (MethodNode, InternalName)]
             (declarationClassNode, source) <- byteCodeRepository.classNodeAndSource(declarationClass): Either[OptimizerWarning, (ClassNode, Source)]
-            declarationClassBType          =  classBTypeFromClassNode(declarationClassNode)
           } yield {
+              val declarationClassBType = classBTypeFromClassNode(declarationClassNode)
               val CallsiteInfo(safeToInline, safeToRewrite, annotatedInline, annotatedNoInline, samParamTypes, warning) = analyzeCallsite(method, declarationClassBType, call.owner, source)
               Callee(
                 callee = method,
@@ -347,6 +349,12 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
   final case class Callsite(callsiteInstruction: MethodInsnNode, callsiteMethod: MethodNode, callsiteClass: ClassBType,
                             callee: Either[OptimizerWarning, Callee], argInfos: IntMap[ArgInfo],
                             callsiteStackHeight: Int, receiverKnownNotNull: Boolean, callsitePosition: Position) {
+    /**
+     * Contains callsites that were created during inlining by cloning this callsite. Used to find
+     * corresponding callsites when inlining post-inline requests.
+     */
+    val inlinedClones = mutable.Set.empty[Callsite]
+
     override def toString =
       "Invocation of" +
         s" ${callee.map(_.calleeDeclarationClass.internalName).getOrElse("?")}.${callsiteInstruction.name + callsiteInstruction.desc}" +
@@ -399,6 +407,10 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
    *                              graph when re-writing a closure invocation to the body method.
    */
   final case class ClosureInstantiation(lambdaMetaFactoryCall: LambdaMetaFactoryCall, ownerMethod: MethodNode, ownerClass: ClassBType, capturedArgInfos: IntMap[ArgInfo]) {
+    /**
+     * Contains closure instantiations that were created during inlining by cloning this instantiation.
+     */
+    val inlinedClones = mutable.Set.empty[ClosureInstantiation]
     override def toString = s"ClosureInstantiation($lambdaMetaFactoryCall, ${ownerMethod.name + ownerMethod.desc}, $ownerClass)"
   }
   final case class LambdaMetaFactoryCall(indy: InvokeDynamicInsnNode, samMethodType: Type, implMethod: Handle, instantiatedMethodType: Type)
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 9a5341e131..4e1ecea217 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/Inliner.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/Inliner.scala
@@ -272,14 +272,12 @@ class Inliner[BT <: BTypes](val btypes: BT) {
   def inline(request: InlineRequest): List[CannotInlineWarning] = canInlineBody(request.callsite) match {
     case Some(w) => List(w)
     case None =>
-      val instructionsMap = inlineCallsite(request.callsite)
+      inlineCallsite(request.callsite)
       val postRequests = request.post.flatMap(post => {
-        // the post-request invocation instruction might not exist anymore: it might have been
-        // inlined itself, or eliminated by DCE.
-        for {
-          inlinedInvocationInstr <- instructionsMap.get(post.callsiteInstruction).map(_.asInstanceOf[MethodInsnNode])
-          inlinedCallsite <- callGraph.callsites(request.callsite.callsiteMethod).get(inlinedInvocationInstr)
-        } yield InlineRequest(inlinedCallsite, post.post)
+        post.callsite.inlinedClones.find(cs => cs.callsiteMethod == request.callsite.callsiteMethod) match {
+          case Some(inlinedPostCallsite) if callGraph.containsCallsite(inlinedPostCallsite) => Some(InlineRequest(inlinedPostCallsite, post.post))
+          case _ => None
+        }
       })
       postRequests flatMap inline
   }
@@ -296,7 +294,7 @@ class Inliner[BT <: BTypes](val btypes: BT) {
    * @return A map associating instruction nodes of the callee with the corresponding cloned
    *         instruction in the callsite method.
    */
-  def inlineCallsite(callsite: Callsite): Map[AbstractInsnNode, AbstractInsnNode] = {
+  def inlineCallsite(callsite: Callsite): Unit = {
     import callsite.{callsiteClass, callsiteMethod, callsiteInstruction, receiverKnownNotNull, callsiteStackHeight}
     val Right(callsiteCallee) = callsite.callee
     import callsiteCallee.{callee, calleeDeclarationClass}
@@ -451,7 +449,7 @@ class Inliner[BT <: BTypes](val btypes: BT) {
     callGraph.callsites(callee).valuesIterator foreach { originalCallsite =>
       val newCallsiteIns = instructionMap(originalCallsite.callsiteInstruction).asInstanceOf[MethodInsnNode]
       val argInfos = originalCallsite.argInfos flatMap mapArgInfo
-      callGraph.addCallsite(Callsite(
+      val newCallsite = Callsite(
         callsiteInstruction = newCallsiteIns,
         callsiteMethod = callsiteMethod,
         callsiteClass = callsiteClass,
@@ -460,19 +458,21 @@ class Inliner[BT <: BTypes](val btypes: BT) {
         callsiteStackHeight = callsiteStackHeight + originalCallsite.callsiteStackHeight,
         receiverKnownNotNull = originalCallsite.receiverKnownNotNull,
         callsitePosition = originalCallsite.callsitePosition
-      ))
+      )
+      originalCallsite.inlinedClones += newCallsite
+      callGraph.addCallsite(newCallsite)
     }
 
     callGraph.closureInstantiations(callee).valuesIterator foreach { originalClosureInit =>
       val newIndy = instructionMap(originalClosureInit.lambdaMetaFactoryCall.indy).asInstanceOf[InvokeDynamicInsnNode]
       val capturedArgInfos = originalClosureInit.capturedArgInfos flatMap mapArgInfo
-      callGraph.addClosureInstantiation(
-        ClosureInstantiation(
-          originalClosureInit.lambdaMetaFactoryCall.copy(indy = newIndy),
-          callsiteMethod,
-          callsiteClass,
-          capturedArgInfos)
-      )
+      val newClosureInit = ClosureInstantiation(
+        originalClosureInit.lambdaMetaFactoryCall.copy(indy = newIndy),
+        callsiteMethod,
+        callsiteClass,
+        capturedArgInfos)
+      originalClosureInit.inlinedClones += newClosureInit
+      callGraph.addClosureInstantiation(newClosureInit)
     }
 
     // Remove the elided invocation from the call graph
@@ -480,8 +480,6 @@ class Inliner[BT <: BTypes](val btypes: BT) {
 
     // Inlining a method body can render some code unreachable, see example above (in runInliner).
     unreachableCodeEliminated -= callsiteMethod
-
-    instructionMap
   }
 
   /**
diff --git a/src/compiler/scala/tools/nsc/backend/jvm/opt/InlinerHeuristics.scala b/src/compiler/scala/tools/nsc/backend/jvm/opt/InlinerHeuristics.scala
index d8f12ffb11..52627e77e6 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/InlinerHeuristics.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/InlinerHeuristics.scala
@@ -19,8 +19,10 @@ class InlinerHeuristics[BT <: BTypes](val bTypes: BT) {
   import inliner._
   import callGraph._
 
-  case class InlineRequest(callsite: Callsite, post: List[PostInlineRequest])
-  case class PostInlineRequest(callsiteInstruction: MethodInsnNode, post: List[PostInlineRequest])
+  case class InlineRequest(callsite: Callsite, post: List[InlineRequest]) {
+    // invariant: all post inline requests denote callsites in the callee of the main callsite
+    for (pr <- post) assert(pr.callsite.callsiteMethod == callsite.callee.get.callee, s"Callsite method mismatch: main $callsite - post ${pr.callsite}")
+  }
 
   /**
    * Select callsites from the call graph that should be inlined, grouped by the containing method.