perf: splice mutated file instead of re-printing entire source tree (#2083)

Author: hugo-vrijswijk

modules/core/src/main/scala/stryker4s/model/MutatedFile.scala

@@ -1,8 +1,49 @@
 package stryker4s.model
 
+import cats.Order
+import cats.data.NonEmptySet
+import cats.effect.IO
+import fs2.Stream
 import fs2.io.file.Path
 import stryker4s.mutants.tree.MutantsWithId
 
-import scala.meta.Tree
+import scala.meta.prettyprinters.XtensionReprint
+import scala.meta.{Term, Tree}
 
-final case class MutatedFile(fileOrigin: Path, mutatedSource: Tree, mutants: MutantsWithId)
+final case class MutatedFile(
+    fileOrigin: Path,
+    mutatedSource: Tree,
+    mutants: MutantsWithId,
+    splice: Option[SourceSplice] = None
+) {
+
+  def mutatedSourceText: Stream[IO, String] =
+    splice.fold(Stream.eval(IO(mutatedSource.text)))(_.render)
+}
+
+/** Segments needed to serialize a mutated file
+  */
+final case class SourceSplice(originalText: String, replacements: NonEmptySet[SourceReplacement]) {
+
+  /** Emits the file in parts:
+    *
+    *   1. each original slice preceding a mutation switch
+    *   2. then the rendered switch
+    *   3. finally the trailing original slice.
+    */
+  def render: Stream[IO, String] = {
+    val switches = Stream.emits(replacements.toNonEmptyList.toList).zipWithPrevious.flatMap { case (prev, r) =>
+      val sliceStart = prev.fold(0)(_.endOffset)
+      Stream.emit(originalText.substring(sliceStart, r.begOffset)) ++ Stream.eval(IO(r.tree.reprint()))
+    }
+    switches ++ Stream.emit(originalText.substring(replacements.last.endOffset))
+  }
+}
+
+/** A mutation switch that replaces the original source in the half-open range `[begOffset, endOffset)`. */
+final case class SourceReplacement(begOffset: Int, endOffset: Int, tree: Term)
+
+object SourceReplacement {
+  implicit val order: Order[SourceReplacement] = Order.by(_.begOffset)
+  implicit val ordering: Ordering[SourceReplacement] = order.toOrdering
+}

modules/core/src/main/scala/stryker4s/mutants/tree/MutantInstrumenter.scala

@@ -7,9 +7,10 @@ import stryker4s.exception.{Stryker4sException, UnableToBuildPatternMatchExcepti
 import stryker4s.extension.TreeExtensions.{treeEq, TransformOnceExtension}
 import stryker4s.log.Logger
 import stryker4s.model.*
+import stryker4s.model.SourceReplacement.*
 
+import scala.collection.immutable.SortedSet
 import scala.meta.*
-import scala.util.control.NonFatal
 import scala.util.{Failure, Success, Try}
 
 /** Instrument (place) mutants in a tree
@@ -22,23 +23,29 @@ class MutantInstrumenter(options: InstrumenterOptions)(implicit log: Logger) {
 
   def instrumentFile(context: SourceContext, mutantMap: Map[PlaceableTree, MutantsWithId]): MutatedFile = {
 
-    def instrumentWithMutants(mutantMap: Map[PlaceableTree, MutantsWithId]): PartialFunction[Tree, Tree] = {
+    // Rendered mutation switches for each outermost mutated statement, used to splice the file together
+    val spliceReplacements = SortedSet.newBuilder[SourceReplacement]
+
+    def instrumentWithMutants(
+        mutantMap: Map[PlaceableTree, MutantsWithId],
+        record: Boolean
+    ): PartialFunction[Tree, Tree] = {
 
       Function.unlift { originalTree =>
         val p = PlaceableTree(originalTree)
         mutantMap.get(p).map { case mutations =>
           val mutableCases = mutations.map(mutantToCase)
 
           // Continue deeper into the tree (without the currently placed mutants)
-          val withDefaultsTransformed = PlaceableTree(p.tree.transformOnce(instrumentWithMutants(mutantMap - p)))
+          val withDefaultsTransformed =
+            PlaceableTree(p.tree.transformOnce(instrumentWithMutants(mutantMap - p, record = false)))
           val default = defaultCase(withDefaultsTransformed, mutations.map(_.id).toNonEmptyList)
 
           val cases = mutableCases :+ default
 
-          try
-            buildMatch(cases)
-          catch {
-            case NonFatal(e) =>
+          val mutationSwitch = Either
+            .catchNonFatal(buildMatch(cases))
+            .valueOr { e =>
               log.error(
                 s"Failed to instrument mutants in `${context.path}`. Original statement: [${originalTree.text}]"
               )
@@ -50,12 +57,17 @@ class MutantInstrumenter(options: InstrumenterOptions)(implicit log: Logger) {
                 e
               )
               throw UnableToBuildPatternMatchException(context.path)
-          }
+            }
+
+          if (record)
+            spliceReplacements += SourceReplacement(p.tree.begOffset, p.tree.endOffset, mutationSwitch)
+
+          mutationSwitch
         }
       }
     }
 
-    val newTree = Try(context.source.transformOnce(instrumentWithMutants(mutantMap))) match {
+    val newTree = Try(context.source.transformOnce(instrumentWithMutants(mutantMap, record = true))) match {
       case Success(tree)                  => tree
       case Failure(e: Stryker4sException) => throw e
       case Failure(e)                     =>
@@ -64,8 +76,9 @@ class MutantInstrumenter(options: InstrumenterOptions)(implicit log: Logger) {
     }
 
     val mutations: MutantsWithId = mutantMap.map(_._2).toVector.toNev.get.flatten
+    val splice = spliceReplacements.result().toNes.map(SourceSplice(context.source.pos.input.text, _))
 
-    MutatedFile(context.path, newTree, mutations)
+    MutatedFile(context.path, newTree, mutations, splice)
   }
 
   def mutantToCase(mutant: MutantWithId): Case = {

modules/core/src/main/scala/stryker4s/run/MutantRunner.scala

@@ -145,8 +145,7 @@ class MutantRunner(
       val targetPath = mutatedFile.fileOrigin.inSubDir(tmpDir)
       IO(log.debug(s"Writing ${mutatedFile.fileOrigin} file to $targetPath")) *>
         Files[IO].createDirectories(targetPath.parent.get) *>
-        Stream
-          .eval(IO(mutatedFile.mutatedSource.text))
+        mutatedFile.mutatedSourceText
           .through(Files[IO].writeUtf8(targetPath))
           .compile
           .drain

modules/core/src/main/scala/stryker4s/run/RollbackHandler.scala

@@ -70,7 +70,10 @@ object RollbackHandler {
             case Both(mutants, results) =>
               // If not all errors were removed, we can't continue (return a Left)
               errorsWithoutIds.toLeft(
-                mutatedFile.copy(mutatedSource = treeWithoutErrors, mutants = mutants.toNev).asRight -> results
+                // Drop the splice info: it describes the pre-rollback tree, so the rewritten tree must be rendered whole
+                mutatedFile
+                  .copy(mutatedSource = treeWithoutErrors, mutants = mutants.toNev, splice = none)
+                  .asRight -> results
               )
             // All mutants were removed
             case Ior.Right(results) => (mutatedFile.fileOrigin.asLeft, results).asRight