Make Trees.transformCtx return the local context if needed

Author: TheElectronWill

compiler/src/dotty/tools/dotc/ast/Trees.scala

@@ -1368,18 +1368,25 @@ object Trees {
     def localCtx(tree: Tree)(using Context): Context
 
     /** The context to use when transforming a tree.
-      * It ensures that the source information is correct.
+      * It ensures that the source is correct, and that the local context is used if
+      * that's necessary for transforming the whole tree.
       * TODO: ensure transform is always called with the correct context as argument
       * @see https://github.com/lampepfl/dotty/pull/13880#discussion_r836395977
       */
     def transformCtx(tree: Tree)(using Context): Context =
-      if tree.source.exists && tree.source != ctx.source
-      then ctx.withSource(tree.source)
-      else ctx
+      val sourced =
+        if tree.source.exists && tree.source != ctx.source
+        then ctx.withSource(tree.source)
+        else ctx
+      tree match
+        case t: (MemberDef | PackageDef | LambdaTypeTree | TermLambdaTypeTree) =>
+          localCtx(t)(using sourced)
+        case _ =>
+          sourced
 
     abstract class TreeMap(val cpy: TreeCopier = inst.cpy) { self =>
       def transform(tree: Tree)(using Context): Tree = {
-        inContext(transformCtx(tree)){
+        inContext(transformCtx(tree)) {
           Stats.record(s"TreeMap.transform/$getClass")
           if (skipTransform(tree)) tree
           else tree match {
@@ -1436,13 +1443,9 @@ object Trees {
             case AppliedTypeTree(tpt, args) =>
               cpy.AppliedTypeTree(tree)(transform(tpt), transform(args))
             case LambdaTypeTree(tparams, body) =>
-              inContext(localCtx(tree)) {
-                cpy.LambdaTypeTree(tree)(transformSub(tparams), transform(body))
-              }
+              cpy.LambdaTypeTree(tree)(transformSub(tparams), transform(body))
             case TermLambdaTypeTree(params, body) =>
-              inContext(localCtx(tree)) {
-                cpy.TermLambdaTypeTree(tree)(transformSub(params), transform(body))
-              }
+              cpy.TermLambdaTypeTree(tree)(transformSub(params), transform(body))
             case MatchTypeTree(bound, selector, cases) =>
               cpy.MatchTypeTree(tree)(transform(bound), transform(selector), transformSub(cases))
             case ByNameTypeTree(result) =>
@@ -1458,30 +1461,21 @@ object Trees {
             case EmptyValDef =>
               tree
             case tree @ ValDef(name, tpt, _) =>
-              inContext(localCtx(tree)) {
-                val tpt1 = transform(tpt)
-                val rhs1 = transform(tree.rhs)
-                cpy.ValDef(tree)(name, tpt1, rhs1)
-              }
+              val tpt1 = transform(tpt)
+              val rhs1 = transform(tree.rhs)
+              cpy.ValDef(tree)(name, tpt1, rhs1)
             case tree @ DefDef(name, paramss, tpt, _) =>
-              inContext(localCtx(tree)) {
-                cpy.DefDef(tree)(name, transformParamss(paramss), transform(tpt), transform(tree.rhs))
-              }
+              cpy.DefDef(tree)(name, transformParamss(paramss), transform(tpt), transform(tree.rhs))
             case tree @ TypeDef(name, rhs) =>
-              inContext(localCtx(tree)) {
-                cpy.TypeDef(tree)(name, transform(rhs))
-              }
+              cpy.TypeDef(tree)(name, transform(rhs))
             case tree @ Template(constr, parents, self, _) if tree.derived.isEmpty =>
               cpy.Template(tree)(transformSub(constr), transform(tree.parents), Nil, transformSub(self), transformStats(tree.body, tree.symbol))
             case Import(expr, selectors) =>
               cpy.Import(tree)(transform(expr), selectors)
             case Export(expr, selectors) =>
               cpy.Export(tree)(transform(expr), selectors)
             case PackageDef(pid, stats) =>
-              val pid1 = transformSub(pid)
-              inContext(localCtx(tree)) {
-                cpy.PackageDef(tree)(pid1, transformStats(stats, ctx.owner))
-              }
+              cpy.PackageDef(tree)(transformSub(pid), transformStats(stats, ctx.owner))
             case Annotated(arg, annot) =>
               cpy.Annotated(tree)(transform(arg), transform(annot))
             case Thicket(trees) =>

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

@@ -140,20 +140,21 @@ class InstrumentCoverage extends MacroTransform with IdentityDenotTransformer:
 
           case tree: CaseDef => instrumentCaseDef(tree)
           case tree: ValDef =>
-            // only transform the rhs, in the local context
-            val rhs = transform(tree.rhs)(using localCtx(tree))
+            // only transform the rhs
+            val rhs = transform(tree.rhs)
             cpy.ValDef(tree)(rhs=rhs)
 
           case tree: DefDef =>
             // only transform the params (for the default values) and the rhs
-            val defCtx = localCtx(tree)
-            val paramss = transformParamss(tree.paramss)(using defCtx)
-            val rhs = transform(tree.rhs)(using defCtx)
+            // force instrumentation of literals and other small trees in the rhs,
+            // to ensure that the method call are recorded
+            val paramss = transformParamss(tree.paramss)
+            val rhs = transform(tree.rhs)
             cpy.DefDef(tree)(tree.name, paramss, tree.tpt, rhs)
 
           case tree: PackageDef =>
             // only transform the statements of the package
-            cpy.PackageDef(tree)(tree.pid, transform(tree.stats)(using localCtx(tree)))
+            cpy.PackageDef(tree)(tree.pid, transform(tree.stats))
           case tree: Assign =>
             // only transform the rhs
             cpy.Assign(tree)(tree.lhs, transform(tree.rhs))