Also allow to count allocations in instrumentation

Author: odersky

compiler/src/dotty/tools/dotc/Compiler.scala

@@ -97,7 +97,7 @@ class Compiler {
     List(new Constructors,           // Collect initialization code in primary constructors
                                         // Note: constructors changes decls in transformTemplate, no InfoTransformers should be added after it
          new FunctionalInterfaces,   // Rewrites closures to implement @specialized types of Functions.
-         new InstrumentClosures,     // Count closure allocations under -Yinstrument-closures
+         new Instrumentation,     // Count closure allocations under -Yinstrument-closures
          new GetClass) ::            // Rewrites getClass calls on primitive types.
     List(new LinkScala2Impls,        // Redirect calls to trait methods defined by Scala 2.x, so that they now go to their implementations
          new LambdaLift,             // Lifts out nested functions to class scope, storing free variables in environments

compiler/src/dotty/tools/dotc/config/ScalaSettings.scala

@@ -150,6 +150,7 @@ class ScalaSettings extends Settings.SettingGroup {
   val YnoDecodeStacktraces: Setting[Boolean] = BooleanSetting("-Yno-decode-stacktraces", "Show raw StackOverflow stacktraces, instead of decoding them into triggering operations.")
 
   val YinstrumentClosures: Setting[Boolean] = BooleanSetting("-Yinstrument-closures", "Add instrumentation code that counts closure creations.")
+  val YinstrumentAllocations: Setting[Boolean] = BooleanSetting("-Yinstrument-allocations", "Add instrumentation code that counts allocations.")
 
   /** Dottydoc specific settings */
   val siteRoot: Setting[String] = StringSetting(

compiler/src/dotty/tools/dotc/transform/Instrumentation.scala

@@ -0,0 +1,60 @@
+package dotty.tools.dotc
+package transform
+
+import core._
+import Contexts.Context
+import Symbols._
+import Flags._
+import SymDenotations._
+
+import Decorators._
+import ast.Trees._
+import MegaPhase._
+import StdNames.nme
+import Names.TermName
+import Constants.Constant
+
+
+/** The phase is enabled if a -Yinstrument-... option is set.
+ *  If enabled, it counts the number of closures or allocations for each source position.
+ *  It does this by generating a call to dotty.tools.dotc.util.Stats.doRecord.
+ */
+class Instrumentation extends MiniPhase { thisPhase =>
+  import ast.tpd._
+
+  override def phaseName: String = "instrumentation"
+
+  override def isEnabled(implicit ctx: Context) =
+    ctx.settings.YinstrumentClosures.value ||
+    ctx.settings.YinstrumentAllocations.value
+
+  private var consName: TermName = _
+  private var consEqName: TermName = _
+
+  override def prepareForUnit(tree: Tree)(implicit ctx: Context): Context = {
+    consName = "::".toTermName
+    consEqName = "+=".toTermName
+    ctx
+  }
+
+  private def record(category: String, tree: Tree)(implicit ctx: Context): Tree = {
+    val key = Literal(Constant(s"$category${tree.sourcePos.show}"))
+    ref(defn.Stats_doRecord).appliedTo(key, Literal(Constant(1)))
+  }
+
+  override def transformApply(tree: Apply)(implicit ctx: Context): Tree = tree.fun match {
+    case Select(nu: New, _) =>
+      cpy.Block(tree)(record(i"alloc/${nu.tpe}@", tree) :: Nil, tree)
+    case Select(_, name) if name == consName || name == consEqName =>
+      cpy.Block(tree)(record("alloc/::", tree) :: Nil, tree)
+    case _ =>
+      tree
+  }
+
+  override def transformBlock(tree: Block)(implicit ctx: Context): Block = tree.expr match {
+    case _: Closure =>
+      cpy.Block(tree)(record("closure/", tree) :: tree.stats, tree.expr)
+    case _ =>
+      tree
+  }
+}