Allow @inline/noinline at callsites (in addition to def-site)

Allow annotating individual callsites @inline / @noinline using an
annotation ascription

    c.foo(): @inline

7 files changed, 52 insertions, 25 deletions

diff --git a/src/compiler/scala/tools/nsc/backend/jvm/BCodeBodyBuilder.scala b/src/compiler/scala/tools/nsc/backend/jvm/BCodeBodyBuilder.scala
index a42332f7f2..c4eb6e1b42 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/BCodeBodyBuilder.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/BCodeBodyBuilder.scala
@@ -16,6 +16,7 @@ import java.lang.invoke.LambdaMetafactory
 import scala.tools.asm
 import GenBCode._
 import BackendReporting._
+import scala.tools.asm.tree.MethodInsnNode
 
 /*
  *
@@ -706,6 +707,21 @@ abstract class BCodeBodyBuilder extends BCodeSkelBuilder {
               }
               else {
                 genCallMethod(sym, invokeStyle, app.pos, hostClass)
+                // Check if the Apply tree has an InlineAnnotatedAttachment, added by the typer
+                // for callsites marked `f(): @inline/noinline`. For nullary calls, the attachment
+                // is on the Select node (not on the Apply node added by UnCurry).
+                def checkInlineAnnotated(t: Tree): Unit = {
+                  if (t.hasAttachment[InlineAnnotatedAttachment]) bc.jmethod.instructions.getLast match {
+                    case m: MethodInsnNode =>
+                      if (app.hasAttachment[NoInlineCallsiteAttachment.type]) noInlineAnnotatedCallsites += m
+                      else inlineAnnotatedCallsites += m
+                    case _ =>
+                  } else t match {
+                    case Apply(fun, _) => checkInlineAnnotated(fun)
+                    case _ =>
+                  }
+                }
+                checkInlineAnnotated(app)
               }
 
             } // end of genNormalMethodCall()
diff --git a/src/compiler/scala/tools/nsc/backend/jvm/BTypes.scala b/src/compiler/scala/tools/nsc/backend/jvm/BTypes.scala
index f1b515910f..92aaf991bf 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/BTypes.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/BTypes.scala
@@ -82,6 +82,14 @@ abstract class BTypes {
   val callsitePositions: concurrent.Map[MethodInsnNode, Position] = recordPerRunCache(TrieMap.empty)
 
   /**
+   * Stores callsite instructions of invocatinos annotated `f(): @inline/noinline`.
+   * Instructions are added during code generation (BCodeBodyBuilder). The maps are then queried
+   * when building the CallGraph, every Callsite object has an annotated(No)Inline field.
+   */
+  val inlineAnnotatedCallsites: mutable.Set[MethodInsnNode] = recordPerRunCache(mutable.Set.empty)
+  val noInlineAnnotatedCallsites: mutable.Set[MethodInsnNode] = recordPerRunCache(mutable.Set.empty)
+
+  /**
    * Contains the internal names of all classes that are defined in Java source files of the current
    * compilation run (mixed compilation). Used for more detailed error reporting.
    */
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 c804c228f5..66810176a1 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/CallGraph.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/CallGraph.scala
@@ -161,7 +161,9 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
             argInfos = argInfos,
             callsiteStackHeight = a.frameAt(call).getStackSize,
             receiverKnownNotNull = receiverNotNull,
-            callsitePosition = callsitePositions.getOrElse(call, NoPosition)
+            callsitePosition = callsitePositions.getOrElse(call, NoPosition),
+            annotatedInline = inlineAnnotatedCallsites(call),
+            annotatedNoInline = noInlineAnnotatedCallsites(call)
           )
 
         case LambdaMetaFactoryCall(indy, samMethodType, implMethod, instantiatedMethodType) if a.frameAt(indy) != null =>
@@ -348,7 +350,8 @@ 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) {
+                            callsiteStackHeight: Int, receiverKnownNotNull: Boolean, callsitePosition: Position,
+                            annotatedInline: Boolean, annotatedNoInline: Boolean) {
     /**
      * Contains callsites that were created during inlining by cloning this callsite. Used to find
      * corresponding callsites when inlining post-inline requests.
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 a7c4c27a97..4203a93f2e 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/ClosureOptimizer.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/ClosureOptimizer.scala
@@ -263,7 +263,9 @@ class ClosureOptimizer[BT <: BTypes](val btypes: BT) {
       argInfos = argInfos,
       callsiteStackHeight = invocationStackHeight,
       receiverKnownNotNull = true, // see below (*)
-      callsitePosition = originalCallsite.map(_.callsitePosition).getOrElse(NoPosition)
+      callsitePosition = originalCallsite.map(_.callsitePosition).getOrElse(NoPosition),
+      annotatedInline = false,
+      annotatedNoInline = false
     )
     // (*) The documentation in class LambdaMetafactory says:
     //     "if implMethod corresponds to an instance method, the first capture argument
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 261b736029..449a56fdd1 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/Inliner.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/Inliner.scala
@@ -154,10 +154,8 @@ class Inliner[BT <: BTypes](val btypes: BT) {
           if (selfParamType.info.get.inlineInfo.sam.isEmpty) samParamTypes - 0
           else samParamTypes.updated(0, selfParamType)
         }
-        val staticCallsite = Callsite(
+        val staticCallsite = callsite.copy(
           callsiteInstruction = newCallsiteInstruction,
-          callsiteMethod      = callsite.callsiteMethod,
-          callsiteClass       = callsite.callsiteClass,
           callee              = Right(Callee(
             callee                 = implClassMethod,
             calleeDeclarationClass = implClassBType,
@@ -166,11 +164,7 @@ class Inliner[BT <: BTypes](val btypes: BT) {
             annotatedInline        = annotatedInline,
             annotatedNoInline      = annotatedNoInline,
             samParamTypes          = staticCallSamParamTypes,
-            calleeInfoWarning      = infoWarning)),
-          argInfos            = callsite.argInfos,
-          callsiteStackHeight = callsite.callsiteStackHeight,
-          receiverKnownNotNull = callsite.receiverKnownNotNull,
-          callsitePosition = callsite.callsitePosition
+            calleeInfoWarning      = infoWarning))
         )
         callGraph.addCallsite(staticCallsite)
       }
@@ -501,15 +495,12 @@ 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
-      val newCallsite = Callsite(
+      val newCallsite = originalCallsite.copy(
         callsiteInstruction = newCallsiteIns,
         callsiteMethod = callsiteMethod,
         callsiteClass = callsiteClass,
-        callee = originalCallsite.callee,
         argInfos = argInfos,
-        callsiteStackHeight = callsiteStackHeight + originalCallsite.callsiteStackHeight,
-        receiverKnownNotNull = originalCallsite.receiverKnownNotNull,
-        callsitePosition = originalCallsite.callsitePosition
+        callsiteStackHeight = callsiteStackHeight + originalCallsite.callsiteStackHeight
       )
       originalCallsite.inlinedClones += ClonedCallsite(newCallsite, callsite)
       callGraph.addCallsite(newCallsite)
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 52627e77e6..89a768fd9c 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/InlinerHeuristics.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/InlinerHeuristics.scala
@@ -41,18 +41,18 @@ class InlinerHeuristics[BT <: BTypes](val bTypes: BT) {
     compilingMethods.map(methodNode => {
       var requests = Set.empty[InlineRequest]
       callGraph.callsites(methodNode).valuesIterator foreach {
-        case callsite @ Callsite(_, _, _, Right(Callee(callee, calleeDeclClass, safeToInline, _, annotatedInline, _, _, callsiteWarning)), _, _, _, pos) =>
+        case callsite @ Callsite(_, _, _, Right(Callee(callee, calleeDeclClass, safeToInline, _, calleeAnnotatedInline, _, _, callsiteWarning)), _, _, _, pos, _, _) =>
           inlineRequest(callsite) match {
             case Some(Right(req)) => requests += req
             case Some(Left(w))    =>
-              if ((annotatedInline && bTypes.compilerSettings.YoptWarningEmitAtInlineFailed) || w.emitWarning(compilerSettings)) {
-                val annotWarn = if (annotatedInline) " is annotated @inline but" else ""
+              if ((calleeAnnotatedInline && bTypes.compilerSettings.YoptWarningEmitAtInlineFailed) || w.emitWarning(compilerSettings)) {
+                val annotWarn = if (calleeAnnotatedInline) " is annotated @inline but" else ""
                 val msg = s"${BackendReporting.methodSignature(calleeDeclClass.internalName, callee)}$annotWarn could not be inlined:\n$w"
                 backendReporting.inlinerWarning(callsite.callsitePosition, msg)
               }
 
             case None =>
-              if (annotatedInline && bTypes.compilerSettings.YoptWarningEmitAtInlineFailed) {
+              if (calleeAnnotatedInline && !callsite.annotatedNoInline && bTypes.compilerSettings.YoptWarningEmitAtInlineFailed) {
                 // if the callsite is annotated @inline, we report an inline warning even if the underlying
                 // reason is, for example, mixed compilation (which has a separate -Yopt-warning flag).
                 def initMsg = s"${BackendReporting.methodSignature(calleeDeclClass.internalName, callee)} is annotated @inline but cannot be inlined"
@@ -69,7 +69,7 @@ class InlinerHeuristics[BT <: BTypes](val bTypes: BT) {
               }
           }
 
-        case Callsite(ins, _, _, Left(warning), _, _, _, pos) =>
+        case Callsite(ins, _, _, Left(warning), _, _, _, pos, _, _) =>
           if (warning.emitWarning(compilerSettings))
             backendReporting.inlinerWarning(pos, s"failed to determine if ${ins.name} should be inlined:\n$warning")
       }
@@ -108,12 +108,11 @@ class InlinerHeuristics[BT <: BTypes](val bTypes: BT) {
         else None
 
       case "default" =>
-        if (callee.safeToInline && !callee.annotatedNoInline) {
-          val shouldInlineHO = callee.samParamTypes.nonEmpty && (callee.samParamTypes exists {
+        if (callee.safeToInline && !callee.annotatedNoInline && !callsite.annotatedNoInline) {
+          def shouldInlineHO = callee.samParamTypes.nonEmpty && (callee.samParamTypes exists {
             case (index, _) => callsite.argInfos.contains(index)
           })
-
-          if (shouldInlineHO || callee.annotatedInline) Some(requestIfCanInline(callsite))
+          if (callee.annotatedInline || callsite.annotatedInline || shouldInlineHO) Some(requestIfCanInline(callsite))
           else None
         } else None
     }
diff --git a/src/compiler/scala/tools/nsc/typechecker/Typers.scala b/src/compiler/scala/tools/nsc/typechecker/Typers.scala
index 26e04edcca..45ebbd532d 100644
--- a/src/compiler/scala/tools/nsc/typechecker/Typers.scala
+++ b/src/compiler/scala/tools/nsc/typechecker/Typers.scala
@@ -4130,6 +4130,14 @@ trait Typers extends Adaptations with Tags with TypersTracking with PatternTyper
             ann setType arg1.tpe.withAnnotation(annotInfo)
           }
           val atype = ann.tpe
+          // For `f(): @inline/noinline` callsites, add the InlineAnnotatedAttachment. TypeApplys
+          // are eliminated by erasure, so add it to the underlying function in this case.
+          def setInlineAttachment(t: Tree, att: InlineAnnotatedAttachment): Unit = t match {
+            case TypeApply(fun, _) => setInlineAttachment(fun, att)
+            case _ => t.updateAttachment(att)
+          }
+          if (atype.hasAnnotation(definitions.ScalaNoInlineClass)) setInlineAttachment(arg1, NoInlineCallsiteAttachment)
+          else if (atype.hasAnnotation(definitions.ScalaInlineClass)) setInlineAttachment(arg1, InlineCallsiteAttachment)
           Typed(arg1, resultingTypeTree(atype)) setPos tree.pos setType atype
         }
       }