Scoverage: instrument chained calls correctly

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

@@ -6,6 +6,7 @@ import java.nio.file.{Files, Path}
 
 import ast.tpd
 import ast.tpd.*
+import ast.desugar.TrailingForMap
 import collection.mutable
 import core.Comments.Comment
 import core.Flags.*
@@ -93,10 +94,24 @@ object LiftCoverage extends LiftImpure:
       case _ if valueType.existsPart(_.typeSymbol == defn.TypeBox_CAP) => valueType
       case _ => super.liftedExprType(expr)
 
-  def liftForCoverage(defs: mutable.ListBuffer[tpd.Tree], tree: tpd.Apply)(using Context) =
-    val liftedFun = liftApp(defs, tree.fun)
-    val liftedArgs = liftArgs(defs, tree.fun.tpe, tree.args)(using liftingArgsContext)
-    tpd.cpy.Apply(tree)(liftedFun, liftedArgs)
+  def liftForCoverage(
+    defs: mutable.ListBuffer[tpd.Tree],
+    tree: tpd.Apply,
+    shouldLiftSelectedApply: tpd.Apply => Boolean = _ => false,
+    coverageCallFor: tpd.Apply => Option[tpd.Apply] = _ => None
+  )(using Context): tpd.Tree =
+    def recur(tree: tpd.Apply, instrumentCurrent: Boolean): tpd.Tree =
+      val liftedFun = tree.fun match
+        case sel @ tpd.Select(app: tpd.Apply, name) if shouldLiftSelectedApply(app) =>
+          liftApp(defs, tpd.cpy.Select(sel)(recur(app, instrumentCurrent = true), name))
+        case _ =>
+          liftApp(defs, tree.fun)
+      val liftedArgs = liftArgs(defs, tree.fun.tpe, tree.args)(using liftingArgsContext)
+      val liftedApp = tpd.cpy.Apply(tree)(liftedFun, liftedArgs)
+      if instrumentCurrent then coverageCallFor(tree).foreach(defs += _)
+      liftedApp
+
+    recur(tree, instrumentCurrent = false)
 
 /** Implements code coverage by inserting calls to scala.runtime.coverage.Invoker
   * ("instruments" the source code).
@@ -350,7 +365,7 @@ class InstrumentCoverage extends MacroTransform with IdentityDenotTransformer:
           // Also, tree.fun can be lifted too.
           // See LiftCoverage for the internal working of this lifting.
           val liftedDefs = mutable.ListBuffer[Tree]()
-          val liftedApp = LiftCoverage.liftForCoverage(liftedDefs, app)
+          val liftedApp = LiftCoverage.liftForCoverage(liftedDefs, app, selectedApplyNeedsLift, coverageCallForSelectedApply)
 
           InstrumentedParts(liftedDefs.toList, coverageCall, liftedApp)
         else
@@ -728,35 +743,63 @@ class InstrumentCoverage extends MacroTransform with IdentityDenotTransformer:
      * should not be changed to {val $x = f(); T($x)}(1) but to {val $x = f(); val $y = 1; T($x)($y)}
      */
     private def needsLift(tree: Apply)(using Context): Boolean =
-      def isShortCircuitedOp(sym: Symbol) =
-        sym == defn.Boolean_&& || sym == defn.Boolean_||
-
-      def isUnliftableFun(fun: Tree) =
-        /*
-         * 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.
-         *
-         * On top of that, the `s`, `f` and `raw` string interpolators are special-cased
-         * by the compiler and will disappear in phase StringInterpolatorOpt, therefore
-         * they shouldn't be lifted.
-         */
-        val sym = fun.symbol
-        sym.exists && (
-          isShortCircuitedOp(sym)
-          || StringInterpolatorOpt.isCompilerIntrinsic(sym)
-          || sym == defn.Object_synchronized
-          || isContextFunctionApply(fun)
-        )
-      end isUnliftableFun
+      def hasSelectedApply(fun: Tree): Boolean = fun match
+        case Select(app: Apply, _) => selectedApplyNeedsLift(app)
+        case TypeApply(fn, _) => hasSelectedApply(fn)
+        case _ => false
 
       val fun = tree.fun
       val nestedApplyNeedsLift = fun match
         case a: Apply => needsLift(a)
         case _ => false
 
+      nestedApplyNeedsLift ||
+      !isUnliftableFun(fun)
+      && (
+        !tree.hasAttachment(TrailingForMap) && hasSelectedApply(fun)
+        || !tree.args.isEmpty && !tree.args.forall(LiftCoverage.noLift)
+      )
+
+    private def selectedApplyNeedsLift(tree: Apply)(using Context): Boolean =
+      !LiftCoverage.isUnsafeAssumeSeparate(tree)
+      && canInstrumentApply(tree)
+      && needsLiftWithoutSelectedApply(tree)
+
+    private def needsLiftWithoutSelectedApply(tree: Apply)(using Context): Boolean =
+      val fun = tree.fun
+      val nestedApplyNeedsLift = fun match
+        case a: Apply => needsLiftWithoutSelectedApply(a)
+        case _ => false
+
       nestedApplyNeedsLift ||
       !isUnliftableFun(fun) && !tree.args.isEmpty && !tree.args.forall(LiftCoverage.noLift)
 
+    private def isShortCircuitedOp(sym: Symbol)(using Context) =
+      sym == defn.Boolean_&& || sym == defn.Boolean_||
+
+    private def isUnliftableFun(fun: Tree)(using Context) =
+      /*
+       * 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.
+       *
+       * On top of that, the `s`, `f` and `raw` string interpolators are special-cased
+       * by the compiler and will disappear in phase StringInterpolatorOpt, therefore
+       * they shouldn't be lifted.
+       */
+      val sym = fun.symbol
+      sym.exists && (
+        isShortCircuitedOp(sym)
+        || StringInterpolatorOpt.isCompilerIntrinsic(sym)
+        || sym == defn.Object_synchronized
+        || isContextFunctionApply(fun)
+      )
+    end isUnliftableFun
+
+    private def coverageCallForSelectedApply(tree: Apply)(using Context): Option[Apply] =
+      Option.when(!LiftCoverage.isUnsafeAssumeSeparate(tree) && canInstrumentApply(tree))(
+        createInvokeCall(tree, tree.sourcePos)
+      )
+
     private def isContextFunctionApply(fun: Tree)(using Context): Boolean =
       fun match
         case Select(prefix, nme.apply) =>

tests/coverage/pos/Constructor.scoverage.check

@@ -92,6 +92,23 @@ covtest
 C
 Class
 covtest.C
+<init>
+108
+120
+10
+toString
+Apply
+false
+0
+false
+x.toString()
+
+5
+Constructor.scala
+covtest
+C
+Class
+covtest.C
 f
 133
 138
@@ -103,7 +120,7 @@ false
 false
 def f
 
-5
+6
 Constructor.scala
 covtest
 C
@@ -120,7 +137,7 @@ false
 false
 1
 
-6
+7
 Constructor.scala
 covtest
 C
@@ -137,7 +154,7 @@ false
 false
 def x
 
-7
+8
 Constructor.scala
 covtest
 C
@@ -154,7 +171,7 @@ false
 false
 f(x)
 
-8
+9
 Constructor.scala
 covtest
 C
@@ -171,7 +188,7 @@ false
 false
 x
 
-9
+10
 Constructor.scala
 covtest
 C
@@ -188,7 +205,7 @@ false
 false
 2
 
-10
+11
 Constructor.scala
 covtest
 C
@@ -205,7 +222,7 @@ false
 false
 def g
 
-11
+12
 Constructor.scala
 covtest
 O
@@ -222,7 +239,7 @@ false
 false
 def g
 
-12
+13
 Constructor.scala
 covtest
 O
@@ -239,7 +256,7 @@ false
 false
 1
 
-13
+14
 Constructor.scala
 covtest
 O
@@ -256,7 +273,7 @@ false
 false
 def y
 
-14
+15
 Constructor.scala
 covtest
 O

tests/coverage/pos/ContextFunctions.scoverage.check

@@ -127,6 +127,23 @@ Imperative
 Class
 covtest.Imperative
 readPerson
+252
+295
+14
+<init>
+Apply
+false
+0
+false
+OnError((e) => readName2(using e)(using s))
+
+7
+ContextFunctions.scala
+covtest
+Imperative
+Class
+covtest.Imperative
+readPerson
 192
 206
 12

tests/coverage/pos/DefaultArgs.scoverage.check

@@ -399,6 +399,40 @@ DefaultArgs
 Object
 covtest.DefaultArgs
 staticCaller
+706
+720
+28
+staticMethod
+Apply
+false
+0
+false
+staticMethod()
+
+23
+DefaultArgs.scala
+covtest
+DefaultArgs
+Object
+covtest.DefaultArgs
+staticCaller
+706
+738
+28
++
+Apply
+false
+0
+false
+staticMethod() + staticMethod(5)
+
+24
+DefaultArgs.scala
+covtest
+DefaultArgs
+Object
+covtest.DefaultArgs
+staticCaller
 676
 692
 27