Store information about function literals in call graph

Remember in the call graph if a function literal is passed as an
argument to a higher-order function.

5 files changed, 135 insertions, 41 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 fcb8991baa..bf19bece8b 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/CallGraph.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/CallGraph.scala
@@ -16,7 +16,7 @@ import scala.collection.{concurrent, mutable}
 import scala.collection.convert.decorateAsScala._
 import scala.tools.nsc.backend.jvm.BTypes.InternalName
 import scala.tools.nsc.backend.jvm.BackendReporting._
-import scala.tools.nsc.backend.jvm.analysis.{NotNull, NullnessAnalyzer}
+import scala.tools.nsc.backend.jvm.analysis.{ParameterProducer, ProdConsAnalyzer, NotNull, NullnessAnalyzer}
 import ByteCodeRepository.{Source, CompilationUnit}
 import BytecodeUtils._
 
@@ -114,6 +114,9 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
     var methodCallsites = Map.empty[MethodInsnNode, Callsite]
     var methodClosureInstantiations = Map.empty[InvokeDynamicInsnNode, ClosureInstantiation]
 
+    // lazy so it is only computed if actually used by computeArgInfos
+    lazy val prodCons = new ProdConsAnalyzer(methodNode, definingClass.internalName)
+
     methodNode.instructions.iterator.asScala foreach {
       case call: MethodInsnNode =>
         val callee: Either[OptimizerWarning, Callee] = for {
@@ -133,13 +136,7 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
             calleeInfoWarning = warning)
         }
 
-        val argInfos = if (callee.isLeft) Nil else {
-          // TODO: for now it's Nil, because we don't run any data flow analysis
-          // there's no point in using the parameter types, that doesn't add any information.
-          // NOTE: need to run the same analyses after inlining, to re-compute the argInfos for the
-          // new duplicated callsites, see Inliner.inline
-          Nil
-        }
+        val argInfos = computeArgInfos(callee, call, methodNode, definingClass, Some(prodCons))
 
         val receiverNotNull = call.getOpcode == Opcodes.INVOKESTATIC || {
           val numArgs = Type.getArgumentTypes(call.desc).length
@@ -169,6 +166,50 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
     callsites(methodNode) = methodCallsites
     closureInstantiations(methodNode) = methodClosureInstantiations
   }
+  
+  def computeArgInfos(
+                       callee: Either[OptimizerWarning, Callee],
+                       callsiteInsn: MethodInsnNode, callsiteMethod: MethodNode, callsiteClass: ClassBType,
+                       methodProdCons: => Option[ProdConsAnalyzer] = None): IntMap[ArgInfo] = {
+    if (callee.isLeft) IntMap.empty
+    else {
+      if (callee.get.higherOrderParams.nonEmpty) {
+
+        val prodCons = methodProdCons.getOrElse({
+          localOpt.minimalRemoveUnreachableCode(callsiteMethod, callsiteClass.internalName)
+          new ProdConsAnalyzer(callsiteMethod, callsiteClass.internalName)
+        })
+
+        // TODO: use type analysis instead - should be more efficient than prodCons
+        // some random thoughts:
+        //  - assign special types to parameters and indy-lambda-functions to track them
+        //  - upcast should not change type flow analysis: don't lose information.
+        //  - can we do something about factory calls? Foo(x) for case class foo gives a Foo.
+        //    inline the factory? analysis across method boundry?
+
+        lazy val callFrame = prodCons.frameAt(callsiteInsn)
+        val receiverOrFirstArgSlot = {
+          val numArgs = Type.getArgumentTypes(callsiteInsn.desc).length + (if (callsiteInsn.getOpcode == Opcodes.INVOKESTATIC) 0 else 1)
+          callFrame.stackTop - numArgs + 1
+        }
+        callee.get.higherOrderParams flatMap {
+          case (index, paramType) =>
+            val prods = prodCons.initialProducersForValueAt(callsiteInsn, receiverOrFirstArgSlot + index)
+            if (prods.size != 1) None
+            else {
+              val argInfo = prods.head match {
+                case LambdaMetaFactoryCall(_, _, _, _) => Some(FunctionLiteral)
+                case ParameterProducer(local)          => Some(ForwardedParam(local))
+                case _                                 => None
+              }
+              argInfo.map((index, _))
+            }
+        }
+      } else {
+        IntMap.empty
+      }
+    }
+  }
 
   /**
    * Just a named tuple used as return type of `analyzeCallsite`.
@@ -254,7 +295,7 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
    * @param callsitePosition    The source position of the callsite, used for inliner warnings.
    */
   final case class Callsite(callsiteInstruction: MethodInsnNode, callsiteMethod: MethodNode, callsiteClass: ClassBType,
-                            callee: Either[OptimizerWarning, Callee], argInfos: List[ArgInfo],
+                            callee: Either[OptimizerWarning, Callee], argInfos: IntMap[ArgInfo],
                             callsiteStackHeight: Int, receiverKnownNotNull: Boolean, callsitePosition: Position) {
     override def toString =
       "Invocation of" +
@@ -267,8 +308,9 @@ class CallGraph[BT <: BTypes](val btypes: BT) {
    * Information about invocation arguments, obtained through data flow analysis of the callsite method.
    */
   sealed trait ArgInfo
-  final case class ArgTypeInfo(argType: BType, isPrecise: Boolean, knownNotNull: Boolean) extends ArgInfo
+  case object FunctionLiteral extends ArgInfo
   final case class ForwardedParam(index: Int) extends ArgInfo
+  //  final case class ArgTypeInfo(argType: BType, isPrecise: Boolean, knownNotNull: Boolean) extends ArgInfo
   // can be extended, e.g., with constant types
 
   /**
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 ba6897b06e..70e203aac1 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/ClosureOptimizer.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/ClosureOptimizer.scala
@@ -9,6 +9,7 @@ package opt
 
 import scala.annotation.switch
 import scala.collection.immutable
+import scala.collection.immutable.IntMap
 import scala.reflect.internal.util.NoPosition
 import scala.tools.asm.{Type, Opcodes}
 import scala.tools.asm.tree._
@@ -235,24 +236,25 @@ class ClosureOptimizer[BT <: BTypes](val btypes: BT) {
     // 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 callee = bodyMethod.map({
+      case (bodyMethodNode, bodyMethodDeclClass) =>
+        val bodyDeclClassType = classBTypeFromParsedClassfile(bodyMethodDeclClass)
+        Callee(
+          callee = bodyMethodNode,
+          calleeDeclarationClass = bodyDeclClassType,
+          safeToInline = compilerSettings.YoptInlineGlobal || bodyMethodIsBeingCompiled,
+          safeToRewrite = false, // the lambda body method is not a trait interface method
+          annotatedInline = false,
+          annotatedNoInline = false,
+          inliner.heuristics.higherOrderParams(bodyMethodNode, bodyDeclClassType),
+          calleeInfoWarning = None)
+    })
     val bodyMethodCallsite = Callsite(
       callsiteInstruction = bodyInvocation,
       callsiteMethod = ownerMethod,
       callsiteClass = closureInit.ownerClass,
-      callee = bodyMethod.map({
-        case (bodyMethodNode, bodyMethodDeclClass) =>
-          val bodyDeclClassType = classBTypeFromParsedClassfile(bodyMethodDeclClass)
-          Callee(
-            callee = bodyMethodNode,
-            calleeDeclarationClass = bodyDeclClassType,
-            safeToInline = compilerSettings.YoptInlineGlobal || bodyMethodIsBeingCompiled,
-            safeToRewrite = false, // the lambda body method is not a trait interface method
-            annotatedInline = false,
-            annotatedNoInline = false,
-            inliner.heuristics.higherOrderParams(bodyMethodNode, bodyDeclClassType),
-            calleeInfoWarning = None)
-      }),
-      argInfos = Nil,
+      callee = callee,
+      argInfos = computeArgInfos(callee, bodyInvocation, ownerMethod, closureInit.ownerClass),
       callsiteStackHeight = invocationStackHeight,
       receiverKnownNotNull = true, // see below (*)
       callsitePosition = originalCallsite.map(_.callsitePosition).getOrElse(NoPosition)
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 265685ad84..2918c33c74 100644
--- a/src/compiler/scala/tools/nsc/backend/jvm/opt/Inliner.scala
+++ b/src/compiler/scala/tools/nsc/backend/jvm/opt/Inliner.scala
@@ -8,6 +8,7 @@ package backend.jvm
 package opt
 
 import scala.annotation.tailrec
+import scala.collection.immutable.IntMap
 import scala.tools.asm
 import asm.Handle
 import asm.Opcodes._
@@ -177,7 +178,7 @@ class Inliner[BT <: BTypes](val btypes: BT) {
             annotatedNoInline      = annotatedNoInline,
             higherOrderParams      = staticCallHigherOrderParams,
             calleeInfoWarning      = infoWarning)),
-          argInfos            = Nil,
+          argInfos            = callsite.argInfos,
           callsiteStackHeight = callsite.callsiteStackHeight,
           receiverKnownNotNull = callsite.receiverKnownNotNull,
           callsitePosition = callsite.callsitePosition
@@ -410,6 +411,10 @@ class Inliner[BT <: BTypes](val btypes: BT) {
     callsiteMethod.localVariables.addAll(cloneLocalVariableNodes(callee, labelsMap, callee.name + "_").asJava)
     callsiteMethod.tryCatchBlocks.addAll(cloneTryCatchBlockNodes(callee, labelsMap).asJava)
 
+    callsiteMethod.maxLocals += returnType.getSize + callee.maxLocals
+    val numStoredArgs = calleeParamTypes.length + (if (isStaticMethod(callee)) 0 else 1)
+    callsiteMethod.maxStack = math.max(callsiteMethod.maxStack, callee.maxStack + callsiteStackHeight - numStoredArgs)
+
     callGraph.addIfMissing(callee, calleeDeclarationClass)
 
     // Add all invocation instructions and closure instantiations that were inlined to the call graph
@@ -420,12 +425,11 @@ class Inliner[BT <: BTypes](val btypes: BT) {
         callsiteMethod = callsiteMethod,
         callsiteClass = callsiteClass,
         callee = originalCallsite.callee,
-        argInfos = Nil, // TODO: re-compute argInfos for new destination (once we actually compute them)
+        argInfos = computeArgInfos(originalCallsite.callee, newCallsiteIns, callsiteMethod, callsiteClass), // TODO: try to re-build argInfos from the original callsite's
         callsiteStackHeight = callsiteStackHeight + originalCallsite.callsiteStackHeight,
         receiverKnownNotNull = originalCallsite.receiverKnownNotNull,
         callsitePosition = originalCallsite.callsitePosition
-      )
-      )
+      ))
     }
 
     callGraph.closureInstantiations(callee).valuesIterator foreach { originalClosureInit =>
@@ -441,10 +445,6 @@ class Inliner[BT <: BTypes](val btypes: BT) {
     // Inlining a method body can render some code unreachable, see example above (in runInliner).
     unreachableCodeEliminated -= callsiteMethod
 
-    callsiteMethod.maxLocals += returnType.getSize + callee.maxLocals
-    val numStoredArgs = calleeParamTypes.length + (if (isStaticMethod(callee)) 0 else 1)
-    callsiteMethod.maxStack = math.max(callsiteMethod.maxStack, callee.maxStack + callsiteStackHeight - numStoredArgs)
-
     instructionMap
   }
 
diff --git a/test/junit/scala/tools/nsc/backend/jvm/opt/CallGraphTest.scala b/test/junit/scala/tools/nsc/backend/jvm/opt/CallGraphTest.scala
index 715db3f8c2..b518cbdc50 100644
--- a/test/junit/scala/tools/nsc/backend/jvm/opt/CallGraphTest.scala
+++ b/test/junit/scala/tools/nsc/backend/jvm/opt/CallGraphTest.scala
@@ -6,6 +6,7 @@ import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
 import org.junit.Test
 import scala.collection.generic.Clearable
+import scala.collection.immutable.IntMap
 import scala.tools.asm.Opcodes._
 import org.junit.Assert._
 
@@ -21,18 +22,31 @@ import AsmUtils._
 import BackendReporting._
 
 import scala.collection.convert.decorateAsScala._
+import scala.tools.testing.ClearAfterClass
 
-@RunWith(classOf[JUnit4])
-class CallGraphTest {
-  val compiler = newCompiler(extraArgs = "-Ybackend:GenBCode -Yopt:inline-global -Yopt-warnings")
-  import compiler.genBCode.bTypes._
+object CallGraphTest extends ClearAfterClass.Clearable {
+  var compiler = newCompiler(extraArgs = "-Ybackend:GenBCode -Yopt:inline-global -Yopt-warnings")
+  def clear(): Unit = { compiler = null }
 
   // allows inspecting the caches after a compilation run
-  val notPerRun: List[Clearable] = List(classBTypeFromInternalName, byteCodeRepository.compilingClasses, byteCodeRepository.parsedClasses, callGraph.callsites)
+  val notPerRun: List[Clearable] = List(
+    compiler.genBCode.bTypes.classBTypeFromInternalName,
+    compiler.genBCode.bTypes.byteCodeRepository.compilingClasses,
+    compiler.genBCode.bTypes.byteCodeRepository.parsedClasses,
+    compiler.genBCode.bTypes.callGraph.callsites)
   notPerRun foreach compiler.perRunCaches.unrecordCache
+}
+
+@RunWith(classOf[JUnit4])
+class CallGraphTest extends ClearAfterClass {
+  ClearAfterClass.stateToClear = CallGraphTest
 
-  def compile(code: String, allowMessage: StoreReporter#Info => Boolean): List[ClassNode] = {
-    notPerRun.foreach(_.clear())
+  val compiler = CallGraphTest.compiler
+  import compiler.genBCode.bTypes._
+  import callGraph._
+
+  def compile(code: String, allowMessage: StoreReporter#Info => Boolean = _ => false): List[ClassNode] = {
+    CallGraphTest.notPerRun.foreach(_.clear())
     compileClasses(compiler)(code, allowMessage = allowMessage)
   }
 
@@ -107,7 +121,7 @@ class CallGraphTest {
         assert(callee.annotatedInline == atInline)
         assert(callee.annotatedNoInline == atNoInline)
 
-        assert(callsite.argInfos == List()) // not defined yet
+        assert(callsite.argInfos == IntMap.empty) // no higher-order methods
       } catch {
         case e: Throwable => println(callsite); throw e
       }
@@ -135,4 +149,40 @@ class CallGraphTest {
     checkCallsite(df7Call, t2, cf7, cClassBType, false, true, true)
     checkCallsite(dg1Call, t2, g1, cMClassBType, true, false, false)
   }
+
+  /**
+   * NOTE: if this test fails for you when running within the IDE, it's probably because you're
+   * using 2.12.0-M2 for compilining within the IDE, which doesn't add SAM information to the
+   * InlineInfo attribute. So the InlineInfo in the classfile for Function1 doesn't say that
+   * it's a SAM type. The test passes when running with ant (which does a full bootstrap).
+   */
+  @Test
+  def checkArgInfos(): Unit = {
+    val code =
+      """abstract class C {
+        |  def h(f: Int => Int): Int = f(1)
+        |  def t1 = h(x => x + 1)
+        |  def t2(i: Int, f: Int => Int, z: Int) = h(f) + i - z
+        |  def t3(f: Int => Int) = h(x => f(x + 1))
+        |}
+        |abstract class D {
+        |  def iAmASam(x: Int): Int
+        |  def selfSamCall = iAmASam(10)
+        |}
+        |""".stripMargin
+    val List(c, d) = compile(code)
+
+    def callIn(m: String) = callGraph.callsites.find(_._1.name == m).get._2.values.head
+    val t1h = callIn("t1")
+    assertEquals(t1h.argInfos.toList, List((1, FunctionLiteral)))
+
+    val t2h = callIn("t2")
+    assertEquals(t2h.argInfos.toList, List((1, ForwardedParam(2))))
+
+    val t3h = callIn("t3")
+    assertEquals(t3h.argInfos.toList, List((1, FunctionLiteral)))
+
+    val selfSamCall = callIn("selfSamCall")
+    assertEquals(selfSamCall.argInfos.toList, List((0,ForwardedParam(0))))
+  }
 }
diff --git a/test/junit/scala/tools/nsc/backend/jvm/opt/InlinerTest.scala b/test/junit/scala/tools/nsc/backend/jvm/opt/InlinerTest.scala
index 135ebe9a78..6dfdf20587 100644
--- a/test/junit/scala/tools/nsc/backend/jvm/opt/InlinerTest.scala
+++ b/test/junit/scala/tools/nsc/backend/jvm/opt/InlinerTest.scala
@@ -97,7 +97,7 @@ class InlinerTest extends ClearAfterClass {
       callsiteMethod = callsiteMethod,
       callsiteClass = callsiteClass,
       callee = Right(callGraph.Callee(callee = callee, calleeDeclarationClass = calleeDeclarationClass, safeToInline = true, safeToRewrite = false, annotatedInline = false, annotatedNoInline = false, higherOrderParams = IntMap.empty, calleeInfoWarning = None)),
-      argInfos = Nil,
+      argInfos = IntMap.empty,
       callsiteStackHeight = callsiteStackHeight,
       receiverKnownNotNull = receiverKnownNotNull,
       callsitePosition = NoPosition),