perf: only traverse tree once when building new mutated tree (#2084)

Author: hugo-vrijswijk

modules/core/src/main/scala/stryker4s/extension/TreeExtensions.scala

@@ -80,6 +80,26 @@ object TreeExtensions {
       parentOfType[T](thisTree).isDefined
   }
 
+  implicit final class AncestorsExtension(val thisTree: Tree) extends AnyVal {
+
+    /** The ancestors of this tree, from the immediate parent up to (and including) `root`. This tree itself is not
+      * included.
+      */
+    final def ancestorsUpTo(root: Tree): Seq[Tree] = {
+      val builder = Vector.newBuilder[Tree]
+      @tailrec
+      def loop(tree: Tree): Unit = {
+        builder += tree
+        tree.parent match {
+          case Some(parent) if tree ne root => loop(parent)
+          case _                            => ()
+        }
+      }
+      thisTree.parent.foreach(loop)
+      builder.result()
+    }
+  }
+
   implicit final class GetMods(val tree: Tree) extends AnyVal {
     final def getMods: List[Mod] =
       tree match {

modules/core/src/main/scala/stryker4s/mutants/findmutants/MutantMatcher.scala

@@ -5,7 +5,7 @@ import cats.syntax.all.*
 import mutationtesting.cats.*
 import stryker4s.config.{Config, ExcludedMutation}
 import stryker4s.extension.PartialFunctionOps.*
-import stryker4s.extension.TreeExtensions.{treeEq, PositionExtension, TransformOnceExtension}
+import stryker4s.extension.TreeExtensions.{treeEq, AncestorsExtension, CollectFirstExtension, PositionExtension}
 import stryker4s.model.*
 import stryker4s.mutants.tree.{IgnoredMutation, IgnoredMutations, Mutations}
 import stryker4s.mutation.*
@@ -143,6 +143,18 @@ class MutantMatcherImpl()(implicit config: Config) extends MutantMatcher {
       replacements: NonEmptyVector[T],
       mutationToTerm: T => Term
   ): PlaceableTree => Either[IgnoredMutations, Mutations] = placeableTree => {
+    // Find the node to replace once, so each replacement only rebuilds the path to that node
+    val target = placeableTree.tree
+      .collectFirst {
+        case t if (t eq original) || (t.pos == original.pos && t === original) => t
+      }
+      .getOrElse(
+        throw new RuntimeException(
+          show"Could not transform '${original.text}' in ${placeableTree.tree.text} (${original.pos.toLocation})"
+        )
+      )
+    val pathToTarget = target.ancestorsUpTo(placeableTree.tree)
+
     val mutations = replacements.map { mutations =>
       val tree = mutationToTerm(mutations)
 
@@ -159,18 +171,18 @@ class MutantMatcherImpl()(implicit config: Config) extends MutantMatcher {
         location,
         description
       )
-      val mutatedTopStatement = placeableTree.tree
-        .transformExactlyOnce {
-          case t if (t eq original) || (t.pos == original.pos && t === original) =>
-            tree
-        }
-        .getOrElse(
+      val transformer = new Transformer {
+        override def apply(t: Tree): Tree =
+          if (t eq target) tree
+          else if (pathToTarget.exists(_ eq t)) super.apply(t)
+          else t
+      }
+
+      transformer(placeableTree.tree) match {
+        case t if t eq placeableTree.tree =>
           throw new RuntimeException(
             show"Could not transform '${original.text}' in ${placeableTree.tree.text} (${metadata.location})"
           )
-        )
-
-      mutatedTopStatement match {
         case t: Term => MutatedCode(t, metadata)
         case t       =>
           throw new RuntimeException(

modules/core/src/test/scala/stryker4s/extension/TreeExtensionsTest.scala

@@ -67,6 +67,70 @@ class TreeExtensionsTest extends Stryker4sSuite {
     }
   }
 
+  describe("ancestorsUpTo") {
+    test("returns the parent chain from the immediate parent up to and including root") {
+      val tree = "a + b * c".parseTerm
+      val b = tree.find(Term.Name("b")).value
+
+      val result = b.ancestorsUpTo(tree)
+
+      assert(result.nonEmpty)
+      assert(result.head eq b.parent.value, "should start at the immediate parent")
+      assert(result.last eq tree, "should end at root")
+      // consecutive elements should form a parent chain
+      result.sliding(2).foreach {
+        case Seq(child, parent) => assert(child.parent.value eq parent)
+        case _                  => ()
+      }
+    }
+
+    test("does not include the tree itself") {
+      val tree = "a + b".parseTerm
+      val b = tree.find(Term.Name("b")).value
+
+      val result = b.ancestorsUpTo(tree)
+
+      assert(!result.exists(_ eq b))
+    }
+
+    test("returns only the root when the tree is a direct child of root") {
+      val tree = "a + b".parseTerm
+      val a = tree.find(Term.Name("a")).value
+
+      val result = a.ancestorsUpTo(tree)
+
+      assert(result.loneElement eq tree)
+    }
+
+    test("does not include ancestors above root") {
+      val tree = "def foo = a + b * c".parseDef
+      val mul = tree.find("b * c".parseTerm).value
+      val b = tree.find(Term.Name("b")).value
+
+      val result = b.ancestorsUpTo(mul)
+
+      assert(result.last eq mul, "should stop at root")
+      assert(!result.exists(_ eq tree), "should not walk above root")
+    }
+
+    test("walks up to the top of the tree when root is not an ancestor") {
+      val tree = "a + b * c".parseTerm
+      val unrelated = "x".parseTerm
+      val b = tree.find(Term.Name("b")).value
+
+      val result = b.ancestorsUpTo(unrelated)
+
+      assert(result.head eq b.parent.value)
+      assert(result.last eq tree, "should reach the top of the tree")
+    }
+
+    test("returns empty when the tree has no parent") {
+      val tree = "a + b".parseTerm
+
+      assert(tree.ancestorsUpTo(tree).isEmpty)
+    }
+  }
+
   describe("treeEq") {
     test("equal for same instance") {
       val tree = "a + b".parseTerm