Instrument collection sizes

Allow to record total sizes of collections on which
some selected operation is performed.

# Conflicts:
#	compiler/src/dotty/tools/dotc/transform/Instrumentation.scala

Author: odersky

compiler/src/dotty/tools/dotc/core/Definitions.scala

@@ -845,8 +845,6 @@ class Definitions {
     @tu lazy val Not_value: Symbol = NotClass.companionModule.requiredMethod(nme.value)
 
   @tu lazy val ValueOfClass: ClassSymbol = requiredClass("scala.ValueOf")
-  @tu lazy val StatsModule: Symbol = requiredModule("dotty.tools.dotc.util.Stats")
-    @tu lazy val Stats_doRecord: Symbol = StatsModule.requiredMethod("doRecord")
 
   @tu lazy val FromDigitsClass: ClassSymbol           = requiredClass("scala.util.FromDigits")
   @tu lazy val FromDigits_WithRadixClass: ClassSymbol = requiredClass("scala.util.FromDigits.WithRadix")

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

@@ -1,4 +1,5 @@
-package dotty.tools.dotc
+package dotty.tools
+package dotc
 package transform
 
 import core._
@@ -27,33 +28,57 @@ class Instrumentation extends MiniPhase { thisPhase =>
   override def isEnabled(using Context) =
     ctx.settings.Yinstrument.value
 
-  private val namesOfInterest = List(
-    "::", "+=", "toString", "newArray", "box", "toCharArray",
-    "map", "flatMap", "filter", "withFilter", "collect", "foldLeft", "foldRight", "take",
-    "reverse", "mapConserve", "mapconserve", "filterConserve", "zip",
-    "denotsNamed", "lookup", "lookupEntry", "lookupAll", "toList")
-  private var namesToRecord: Set[Name] = _
+  private val collectionNamesOfInterest = List(
+    "map", "flatMap", "filter", "filterNot", "withFilter", "collect", "flatten", "foldLeft", "foldRight", "take",
+    "reverse", "zip", "++", ":::", ":+", "distinct", "dropRight", "takeRight", "groupBy", "groupMap", "init", "inits",
+    "interect", "mkString", "partition", "reverse_:::", "scanLeft", "scanRight",
+    "sortBy", "sortWith", "sorted", "span", "splitAt", "takeWhile", "transpose", "unzip", "unzip3",
+    "updated", "zipAll", "zipWithIndex",
+    "mapConserve", "mapconserve", "filterConserve", "zipWithConserve", "mapWithIndexConserve"
+  )
+
+  private val namesOfInterest = collectionNamesOfInterest ++ List(
+    "::", "+=", "toString", "newArray", "box", "toCharArray", "termName", "typeName",
+    "slice", "staticRef", "requiredClass")
 
-  private var consName: TermName = _
-  private var consEqName: TermName = _
+  private var namesToRecord: Set[Name] = _
+  private var collectionNamesToRecord: Set[Name] = _
+  private var Stats_doRecord: Symbol = _
+  private var Stats_doRecordSize: Symbol = _
+  private var CollectionIterableClass: ClassSymbol = _
 
   override def prepareForUnit(tree: Tree)(using Context): Context =
     namesToRecord = namesOfInterest.map(_.toTermName).toSet
+    collectionNamesToRecord = collectionNamesOfInterest.map(_.toTermName).toSet
+    val StatsModule = requiredModule("dotty.tools.dotc.util.Stats")
+    Stats_doRecord = StatsModule.requiredMethod("doRecord")
+    Stats_doRecordSize = StatsModule.requiredMethod("doRecordSize")
+    CollectionIterableClass = requiredClass("scala.collection.Iterable")
     ctx
 
   private def record(category: String, tree: Tree)(using Context): Tree = {
     val key = Literal(Constant(s"$category@${tree.sourcePos.show}"))
-    ref(defn.Stats_doRecord).appliedTo(key, Literal(Constant(1)))
+    ref(Stats_doRecord).appliedTo(key, Literal(Constant(1)))
   }
 
+  private def recordSize(tree: Apply)(using Context): Tree = tree.fun match
+    case sel @ Select(qual, name)
+    if collectionNamesToRecord.contains(name)
+       && qual.tpe.widen.derivesFrom(CollectionIterableClass) =>
+      val key = Literal(Constant(s"totalSize/${name} in ${qual.tpe.widen.classSymbol.name}@${tree.sourcePos.show}"))
+      val qual1 = ref(Stats_doRecordSize).appliedTo(key, qual).cast(qual.tpe.widen)
+      cpy.Apply(tree)(cpy.Select(sel)(qual1, name), tree.args)
+    case _ =>
+      tree
+
   private def ok(using Context) =
     !ctx.owner.ownersIterator.exists(_.name.toString.startsWith("Stats"))
 
   override def transformApply(tree: Apply)(using Context): Tree = tree.fun match {
     case Select(nu: New, _) =>
       cpy.Block(tree)(record(i"alloc/${nu.tpe}", tree) :: Nil, tree)
     case ref: RefTree if namesToRecord.contains(ref.name) && ok =>
-      cpy.Block(tree)(record(i"call/${ref.name}", tree) :: Nil, tree)
+      cpy.Block(tree)(record(i"call/${ref.name}", tree) :: Nil, recordSize(tree))
     case _ =>
       tree
   }

compiler/src/dotty/tools/dotc/util/Stats.scala

@@ -28,6 +28,10 @@ import collection.mutable
       hits(name) += n
     }
 
+  def doRecordSize(fn: String, coll: scala.collection.Iterable[_]): coll.type =
+    doRecord(fn, coll.size)
+    coll
+
   inline def trackTime[T](fn: String)(inline op: T): T =
     if (enabled) doTrackTime(fn)(op) else op