|
1 | | -package dotty.tools.dotc |
| 1 | +package dotty.tools |
| 2 | +package dotc |
2 | 3 | package transform |
3 | 4 |
|
4 | 5 | import core._ |
@@ -27,33 +28,57 @@ class Instrumentation extends MiniPhase { thisPhase => |
27 | 28 | override def isEnabled(using Context) = |
28 | 29 | ctx.settings.Yinstrument.value |
29 | 30 |
|
30 | | - private val namesOfInterest = List( |
31 | | - "::", "+=", "toString", "newArray", "box", "toCharArray", |
32 | | - "map", "flatMap", "filter", "withFilter", "collect", "foldLeft", "foldRight", "take", |
33 | | - "reverse", "mapConserve", "mapconserve", "filterConserve", "zip", |
34 | | - "denotsNamed", "lookup", "lookupEntry", "lookupAll", "toList") |
35 | | - private var namesToRecord: Set[Name] = _ |
| 31 | + private val collectionNamesOfInterest = List( |
| 32 | + "map", "flatMap", "filter", "filterNot", "withFilter", "collect", "flatten", "foldLeft", "foldRight", "take", |
| 33 | + "reverse", "zip", "++", ":::", ":+", "distinct", "dropRight", "takeRight", "groupBy", "groupMap", "init", "inits", |
| 34 | + "interect", "mkString", "partition", "reverse_:::", "scanLeft", "scanRight", |
| 35 | + "sortBy", "sortWith", "sorted", "span", "splitAt", "takeWhile", "transpose", "unzip", "unzip3", |
| 36 | + "updated", "zipAll", "zipWithIndex", |
| 37 | + "mapConserve", "mapconserve", "filterConserve", "zipWithConserve", "mapWithIndexConserve" |
| 38 | + ) |
| 39 | + |
| 40 | + private val namesOfInterest = collectionNamesOfInterest ++ List( |
| 41 | + "::", "+=", "toString", "newArray", "box", "toCharArray", "termName", "typeName", |
| 42 | + "slice", "staticRef", "requiredClass") |
36 | 43 |
|
37 | | - private var consName: TermName = _ |
38 | | - private var consEqName: TermName = _ |
| 44 | + private var namesToRecord: Set[Name] = _ |
| 45 | + private var collectionNamesToRecord: Set[Name] = _ |
| 46 | + private var Stats_doRecord: Symbol = _ |
| 47 | + private var Stats_doRecordSize: Symbol = _ |
| 48 | + private var CollectionIterableClass: ClassSymbol = _ |
39 | 49 |
|
40 | 50 | override def prepareForUnit(tree: Tree)(using Context): Context = |
41 | 51 | namesToRecord = namesOfInterest.map(_.toTermName).toSet |
| 52 | + collectionNamesToRecord = collectionNamesOfInterest.map(_.toTermName).toSet |
| 53 | + val StatsModule = requiredModule("dotty.tools.dotc.util.Stats") |
| 54 | + Stats_doRecord = StatsModule.requiredMethod("doRecord") |
| 55 | + Stats_doRecordSize = StatsModule.requiredMethod("doRecordSize") |
| 56 | + CollectionIterableClass = requiredClass("scala.collection.Iterable") |
42 | 57 | ctx |
43 | 58 |
|
44 | 59 | private def record(category: String, tree: Tree)(using Context): Tree = { |
45 | 60 | val key = Literal(Constant(s"$category@${tree.sourcePos.show}")) |
46 | | - ref(defn.Stats_doRecord).appliedTo(key, Literal(Constant(1))) |
| 61 | + ref(Stats_doRecord).appliedTo(key, Literal(Constant(1))) |
47 | 62 | } |
48 | 63 |
|
| 64 | + private def recordSize(tree: Apply)(using Context): Tree = tree.fun match |
| 65 | + case sel @ Select(qual, name) |
| 66 | + if collectionNamesToRecord.contains(name) |
| 67 | + && qual.tpe.widen.derivesFrom(CollectionIterableClass) => |
| 68 | + val key = Literal(Constant(s"totalSize/${name} in ${qual.tpe.widen.classSymbol.name}@${tree.sourcePos.show}")) |
| 69 | + val qual1 = ref(Stats_doRecordSize).appliedTo(key, qual).cast(qual.tpe.widen) |
| 70 | + cpy.Apply(tree)(cpy.Select(sel)(qual1, name), tree.args) |
| 71 | + case _ => |
| 72 | + tree |
| 73 | + |
49 | 74 | private def ok(using Context) = |
50 | 75 | !ctx.owner.ownersIterator.exists(_.name.toString.startsWith("Stats")) |
51 | 76 |
|
52 | 77 | override def transformApply(tree: Apply)(using Context): Tree = tree.fun match { |
53 | 78 | case Select(nu: New, _) => |
54 | 79 | cpy.Block(tree)(record(i"alloc/${nu.tpe}", tree) :: Nil, tree) |
55 | 80 | case ref: RefTree if namesToRecord.contains(ref.name) && ok => |
56 | | - cpy.Block(tree)(record(i"call/${ref.name}", tree) :: Nil, tree) |
| 81 | + cpy.Block(tree)(record(i"call/${ref.name}", tree) :: Nil, recordSize(tree)) |
57 | 82 | case _ => |
58 | 83 | tree |
59 | 84 | } |
|
0 commit comments