Refresh InstrumentCoverage to address review comments

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

@@ -64,11 +64,17 @@ class InstrumentCoverage extends MacroTransform with IdentityDenotTransformer:
 
   override protected def newTransformer(using Context) = CoverageTransormer()
 
-  class CoverageTransormer extends Transformer:
-    var instrumented = false
+  private class CoverageTransormer extends Transformer:
 
     override def transform(tree: Tree)(using Context): Tree =
+      println(tree.show + tree.toString)
       tree match
+        // simple cases
+        case tree: (Literal | Import | Export) => tree
+        case tree: (New | This | Super) => instrument(tree)
+        case tree if (tree.isEmpty || tree.isType) => tree // empty Thicket, Ident, TypTree, ...
+
+        // branches
         case tree: If =>
           cpy.If(tree)(
             cond = transform(tree.cond),
@@ -78,61 +84,51 @@ class InstrumentCoverage extends MacroTransform with IdentityDenotTransformer:
         case tree: Try =>
           cpy.Try(tree)(
             expr = instrument(transform(tree.expr), branch = true),
-            cases = instrumentCasees(tree.cases),
+            cases = instrumentCases(tree.cases),
             finalizer = instrument(transform(tree.finalizer), true)
           )
-        case Apply(fun, _)
-            if (
-              fun.symbol.exists &&
-                fun.symbol.isInstanceOf[Symbol] &&
-                fun.symbol == defn.Boolean_&& || fun.symbol == defn.Boolean_||
-            ) =>
-          super.transform(tree)
-        case tree @ Apply(fun, args) if (fun.isInstanceOf[Apply]) =>
-          // We have nested apply, we have to lift all arguments
-          // Example: def T(x:Int)(y:Int)
-          // T(f())(1) // should not be changed to {val $x = f(); T($x)}(1) but to {val $x = f(); val $y = 1; T($x)($y)}
-          liftApply(tree)
+
+        // f(args)
         case tree: Apply =>
-          if (LiftCoverage.needsLift(tree)) {
+          if needsLift(tree) then
             liftApply(tree)
-          } else {
+          else
             super.transform(tree)
-          }
+
+        // (f(x))[args]
+        case tree @ TypeApply(fun: Apply, args) =>
+          cpy.TypeApply(tree)(transform(fun), args)
+
+        // a.b
         case Select(qual, _) if (qual.symbol.exists && qual.symbol.is(JavaDefined)) =>
           //Java class can't be used as a value, we can't instrument the
           //qualifier ({<Probe>;System}.xyz() is not possible !) instrument it
           //as it is
           instrument(tree)
         case tree: Select =>
-          if (tree.qualifier.isInstanceOf[New]) {
+          if tree.qualifier.isInstanceOf[New] then
             instrument(tree)
-          } else {
+          else
             cpy.Select(tree)(transform(tree.qualifier), tree.name)
-          }
+
         case tree: CaseDef => instrumentCaseDef(tree)
+        case tree: ValDef =>
+          // only transform the rhs
+          cpy.ValDef(tree)(rhs = transform(tree.rhs))
 
-        case tree: Literal => instrument(tree)
-        case tree: Ident if (isWildcardArg(tree)) =>
-          // We don't want to instrument wildcard arguments. `var a = _` can't be instrumented
-          tree
-        case tree: New => instrument(tree)
-        case tree: This => instrument(tree)
-        case tree: Super => instrument(tree)
         case tree: PackageDef =>
-          // We don't instrument the pid of the package, but we do instrument the statements
+          // only transform the statements of the package
           cpy.PackageDef(tree)(tree.pid, transform(tree.stats))
-        case tree: Assign => cpy.Assign(tree)(tree.lhs, transform(tree.rhs))
+        case tree: Assign =>
+          cpy.Assign(tree)(tree.lhs, transform(tree.rhs))
         case tree: Template =>
           // Don't instrument the parents (extends) of a template since it
           // causes problems if the parent constructor takes parameters
           cpy.Template(tree)(
             constr = super.transformSub(tree.constr),
             body = transform(tree.body)
           )
-        case tree: Import => tree
-        // Catch EmptyTree since we can't match directly on it
-        case tree: Thicket if tree.isEmpty => tree
+
         // For everything else just recurse and transform
         case _ =>
           report.warning(
@@ -146,7 +142,7 @@ class InstrumentCoverage extends MacroTransform with IdentityDenotTransformer:
       // NOTE: that if only one arg needs to be lifted, we just lift everything
       val lifted = LiftCoverage.liftForCoverage(buffer, tree)
       val instrumented = buffer.toList.map(transform)
-      //We can now instrument the apply as it is with a custom position to point to the function
+      // We can now instrument the apply as it is with a custom position to point to the function
       Block(
         instrumented,
         instrument(
@@ -156,7 +152,7 @@ class InstrumentCoverage extends MacroTransform with IdentityDenotTransformer:
         )
       )
 
-    def instrumentCasees(cases: List[CaseDef])(using Context): List[CaseDef] =
+    def instrumentCases(cases: List[CaseDef])(using Context): List[CaseDef] =
       cases.map(instrumentCaseDef)
 
     def instrumentCaseDef(tree: CaseDef)(using Context): CaseDef =
@@ -166,7 +162,7 @@ class InstrumentCoverage extends MacroTransform with IdentityDenotTransformer:
       instrument(tree, tree.sourcePos, branch)
 
     def instrument(tree: Tree, pos: SourcePosition, branch: Boolean)(using ctx: Context): Tree =
-      if (pos.exists && !pos.span.isZeroExtent && !tree.isType)
+      if pos.exists && !pos.span.isZeroExtent then
         val id = statementId.incrementAndGet()
         val statement = new Statement(
           source = ctx.source.file.name,
@@ -192,6 +188,30 @@ class InstrumentCoverage extends MacroTransform with IdentityDenotTransformer:
           List(Literal(Constant(id)), Literal(Constant(outputPath)))
         )
 
+    /**
+     * Checks if the apply needs a lift in the coverage phase.
+     * In case of a nested application, we have to lift all arguments
+     * Example:
+     * ```
+     * def T(x:Int)(y:Int)
+     * T(f())(1)
+     * ```
+     * should not be changed to {val $x = f(); T($x)}(1) but to {val $x = f(); val $y = 1; T($x)($y)}
+     */
+    def needsLift(tree: Apply)(using Context): Boolean =
+      // We don't want to lift a || getB(), to avoid calling getB if a is true.
+      // Same idea with a && getB(): if a is false, getB shouldn't be called.
+      def isBooleanOperator(fun: Tree) =
+        fun.symbol.exists &&
+        fun.symbol.isInstanceOf[Symbol] &&
+        fun.symbol == defn.Boolean_&& || fun.symbol == defn.Boolean_||
+
+      val fun = tree.fun
+
+      fun.isInstanceOf[Apply] || // nested apply
+      !isBooleanOperator(fun) ||
+      !tree.args.isEmpty && !tree.args.forall(LiftCoverage.noLift)
+
 object InstrumentCoverage:
   val name: String = "instrumentCoverage"
   val description: String = "instrument code for coverage cheking"

compiler/src/dotty/tools/dotc/typer/EtaExpansion.scala

@@ -11,6 +11,7 @@ import Symbols._
 import Names._
 import NameKinds.UniqueName
 import util.Spans._
+import util.Property
 import collection.mutable
 import Trees._
 
@@ -158,23 +159,20 @@ object LiftComplex extends LiftComplex
 /** Lift complex + lift the prefixes */
 object LiftCoverage extends LiftComplex {
 
-  private var liftEverything = false
+  private val LiftEverything = new Property.Key[Boolean]
 
-  /** Return true if the apply needs a lift in the coverage phase
-    Return false if the args are empty, if one or more will be lifter by a
-    complex lifter.
-   */
-  def needsLift(tree: tpd.Apply)(using Context): Boolean =
-    !tree.args.isEmpty && !tree.args.forall(super.noLift(_))
+  private def liftEverything(using Context): Boolean =
+    ctx.property(LiftEverything).contains(true)
+
+  private def liftEverythingContext(using Context): Context =
+    ctx.fresh.setProperty(LiftEverything, true)
 
   override def noLift(expr: tpd.Tree)(using Context) =
     !liftEverything && super.noLift(expr)
 
   def liftForCoverage(defs: mutable.ListBuffer[tpd.Tree], tree: tpd.Apply)(using Context) = {
     val liftedFun = liftApp(defs, tree.fun)
-    liftEverything = true
-    val liftedArgs = liftArgs(defs, tree.fun.tpe, tree.args)
-    liftEverything = false
+    val liftedArgs = liftArgs(defs, tree.fun.tpe, tree.args)(using liftEverythingContext)
     tpd.cpy.Apply(tree)(liftedFun, liftedArgs)
   }
 }