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

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

    c.foo(): @inline

Author: lrytz

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()

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

@@ -81,6 +81,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.

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.

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

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)

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
     }

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
         }
       }

src/library/scala/inline.scala

@@ -11,8 +11,30 @@
 package scala
 
 /**
- * An annotation on methods that requests that the compiler should
- * try especially hard to inline the annotated method.
+ * An annotation on methods that requests that the compiler should try especially hard to inline the
+ * annotated method. The annotation can be used at definition site or at callsite.
+ *
+ * {{{
+ * @inline   final def f1(x: Int) = x
+ * @noinline final def f2(x: Int) = x
+ *           final def f3(x: Int) = x
+ *
+ *  def t1 = f1(1)              // inlined if possible
+ *  def t2 = f2(1)              // not inlined
+ *  def t3 = f3(1)              // may be inlined (heuristics)
+ *  def t4 = f1(1): @noinline   // not inlined (override at callsite)
+ *  def t5 = f2(1): @inline     // not inlined (cannot override the @noinline at f2's definition)
+ *  def t6 = f3(1): @inline     // inlined if possible
+ *  def t7 = f3(1): @noinline   // not inlined
+ * }
+ * }}}
+ *
+ * Note: parentheses are required when annotating a callsite withing a larger expression.
+ *
+ * {{{
+ * def t1 = f1(1) + f1(1): @noinline   // equivalent to (f1(1) + f1(1)): @noinline
+ * def t2 = f1(1) + (f1(1): @noinline) // the second call to f1 is not inlined
+ * }}}
  *
  * @author Lex Spoon
  * @version 1.0, 2007-5-21

src/library/scala/noinline.scala

@@ -11,8 +11,30 @@
 package scala
 
 /**
- * An annotation on methods that forbids the compiler to inline the
- * method, no matter how safe the inlining appears to be.
+ * An annotation on methods that forbids the compiler to inline the method, no matter how safe the
+ * inlining appears to be. The annotation can be used at definition site or at callsite.
+ *
+ * {{{
+ * @inline   final def f1(x: Int) = x
+ * @noinline final def f2(x: Int) = x
+ *           final def f3(x: Int) = x
+ *
+ *  def t1 = f1(1)              // inlined if possible
+ *  def t2 = f2(1)              // not inlined
+ *  def t3 = f3(1)              // may be inlined (heuristics)
+ *  def t4 = f1(1): @noinline   // not inlined (override at callsite)
+ *  def t5 = f2(1): @inline     // not inlined (cannot override the @noinline at f2's definition)
+ *  def t6 = f3(1): @inline     // inlined if possible
+ *  def t7 = f3(1): @noinline   // not inlined
+ * }
+ * }}}
+ *
+ * Note: parentheses are required when annotating a callsite withing a larger expression.
+ *
+ * {{{
+ * def t1 = f1(1) + f1(1): @noinline   // equivalent to (f1(1) + f1(1)): @noinline
+ * def t2 = f1(1) + (f1(1): @noinline) // the second call to f1 is not inlined
+ * }}}
  *
  * @author Lex Spoon
  * @version 1.0, 2007-5-21

src/reflect/scala/reflect/internal/StdAttachments.scala

@@ -52,4 +52,8 @@ trait StdAttachments {
 
   /** Untyped list of subpatterns attached to selector dummy. */
   case class SubpatternsAttachment(patterns: List[Tree])
+
+  abstract class InlineAnnotatedAttachment
+  case object NoInlineCallsiteAttachment extends InlineAnnotatedAttachment
+  case object InlineCallsiteAttachment extends InlineAnnotatedAttachment
 }