zinc based test quick command

Author: HollandDM

core/define/src/mill/define/Task.scala

@@ -124,7 +124,12 @@ object Task extends TaskBase {
   inline def Command[T](inline t: Result[T])(implicit
       inline w: W[T],
       inline ctx: mill.define.ModuleCtx
-  ): Command[T] = ${ TaskMacros.commandImpl[T]('t)('w, 'ctx, exclusive = '{ false }) }
+  ): Command[T] = ${ TaskMacros.commandImpl[T]('t)('w, 'ctx, exclusive = '{ false }, '{ false }) }
+
+  inline def Command[T](persistent: Boolean)(inline t: Result[T])(implicit
+      inline w: W[T],
+      inline ctx: mill.define.ModuleCtx
+  ): Command[T] = ${ TaskMacros.commandImpl[T]('t)('w, 'ctx, exclusive = '{ false }, '{ persistent }) }
 
   /**
    * @param exclusive Exclusive commands run serially at the end of an evaluation,
@@ -141,7 +146,7 @@ object Task extends TaskBase {
   class CommandFactory private[mill] (val exclusive: Boolean) {
     inline def apply[T](inline t: Result[T])(implicit
         inline w: W[T],
-        inline ctx: mill.define.ModuleCtx
+        inline ctx: mill.define.Ctx
     ): Command[T] = ${ TaskMacros.commandImpl[T]('t)('w, 'ctx, '{ this.exclusive }) }
   }
 
@@ -394,7 +399,8 @@ class Command[+T](
     val ctx0: mill.define.ModuleCtx,
     val writer: W[?],
     val isPrivate: Option[Boolean],
-    val exclusive: Boolean
+    val exclusive: Boolean,
+    override val persistent: Boolean
 ) extends NamedTask[T] {
 
   override def asCommand: Some[Command[T]] = Some(this)
@@ -540,13 +546,13 @@ private object TaskMacros {
       Quotes
   )(t: Expr[Result[T]])(
       w: Expr[W[T]],
-      ctx: Expr[mill.define.ModuleCtx],
+      ctx: Expr[mill.define.Ctx],
       exclusive: Expr[Boolean]
   ): Expr[Command[T]] = {
     appImpl[Command, T](
       (in, ev) =>
         '{
-          new Command[T]($in, $ev, $ctx, $w, ${ taskIsPrivate() }, exclusive = $exclusive)
+          new Command[T]($in, $ev, $ctx, $w, ${ taskIsPrivate() }, exclusive = $exclusive, persistent = $persistent)
         },
       t
     )

scalalib/api/src/mill/scalalib/api/TransitiveSourceStampResults.scala

@@ -0,0 +1,26 @@
+package mill.scalalib.api
+
+final case class TransitiveSourceStampResults(
+  currentStamps: Map[String, String],
+  previousStamps: Option[Map[String, String]] = None
+) {
+  lazy val changedSources: Set[String] = {
+    previousStamps match {
+      case Some(prevStamps) =>
+        currentStamps.view
+        .flatMap { (source, stamp) =>
+          prevStamps.get(source) match {
+            case None => Some(source) // new source
+            case Some(prevStamp) => Option.when(stamp != prevStamp)(source) // changed source
+          }
+        }
+        .toSet
+      case None => currentStamps.keySet
+    }
+  }
+}
+
+object TransitiveSourceStampResults {
+  implicit val jsonFormatter: upickle.default.ReadWriter[TransitiveSourceStampResults] =
+    upickle.default.macroRW
+}

scalalib/package.mill

@@ -25,7 +25,8 @@ object `package` extends RootModule with build.MillStableScalaModule {
       // (also transitively included by com.eed3si9n.jarjarabrams:jarjar-abrams-core)
       // perhaps the class can be copied here?
       Agg(build.Deps.scalaReflect(scalaVersion()))
-    }
+    } ++ 
+    Agg(build.Deps.zinc)
   }
   def testIvyDeps = super.testIvyDeps() ++ Agg(build.Deps.TestDeps.scalaCheck)
   def testTransitiveDeps =

scalalib/src/mill/scalalib/JavaModule.scala

@@ -29,6 +29,10 @@ import mill.scalalib.bsp.BspModule
 import mill.scalalib.publish.Artifact
 import mill.util.Jvm
 import os.Path
+import mill.testrunner.TestResult
+import mill.scalalib.api.TransitiveSourceStampResults
+import scala.collection.immutable.TreeMap
+import scala.util.Try
 
 /**
  * Core configuration required to compile a single Java compilation target
@@ -103,6 +107,173 @@ trait JavaModule
           case _: ClassNotFoundException => // if we can't find the classes, we certainly are not in a ScalaJSModule
         }
     }
+
+    def testQuick(args: String*): Command[(String, Seq[TestResult])] = Task.Command(persistent = true) {
+      val quicktestFailedClassesLog = Task.dest / "quickTestFailedClasses.json"
+      val invalidatedClassesLog = Task.dest / "invalidatedClasses.json"
+      val failedTestClasses =
+        if (!os.exists(quicktestFailedClassesLog)) {
+          Set.empty[String]
+        } else {
+          Try {
+            upickle.default.read[Seq[String]](os.read.stream(quicktestFailedClassesLog))
+          }.getOrElse(Seq.empty[String]).toSet
+        }
+
+      val transitiveStampsFile = Task.dest / "transitiveStamps.json"
+      val previousStampsOpt = if (os.exists(transitiveStampsFile)) {
+        val previousStamps = upickle.default.read[TransitiveSourceStampResults](
+          os.read.stream(transitiveStampsFile)
+        ).currentStamps
+        os.remove(transitiveStampsFile)
+        Some(previousStamps)
+      } else {
+        None
+      }
+
+      def getAnalysisStore(compileResult: CompilationResult): Option[xsbti.compile.CompileAnalysis] = {
+        val analysisStore = sbt.internal.inc.consistent.ConsistentFileAnalysisStore.binary(
+          file = compileResult.analysisFile.toIO,
+          mappers = xsbti.compile.analysis.ReadWriteMappers.getEmptyMappers(),
+          reproducible = true,
+          parallelism = math.min(Runtime.getRuntime.availableProcessors(), 8)
+        )
+        val analysisOptional = analysisStore.get()
+        if (analysisOptional.isPresent) Some(analysisOptional.get.getAnalysis) else None
+      }
+
+      val combinedAnalysis = (compile() +: upstreamCompileOutput())
+        .flatMap(getAnalysisStore)
+        .flatMap {
+          case analysis: sbt.internal.inc.Analysis => Some(analysis)
+          case _ => None
+        }
+        .foldLeft(sbt.internal.inc.Analysis.empty)(_ ++ _)
+
+      val result = TransitiveSourceStampResults(
+        currentStamps = TreeMap.from(
+          combinedAnalysis.stamps.sources.view.map { (source, stamp) =>
+            source.id() -> stamp.writeStamp()
+          }
+        ),
+        previousStamps = previousStampsOpt
+      )
+
+      def getInvalidatedClasspaths(
+        initialInvalidatedClassNames: Set[String],
+        relations: sbt.internal.inc.Relations
+      ): Set[os.Path] = {
+        val seen = collection.mutable.Set.empty[String]
+        val seenList = collection.mutable.Buffer.empty[String]
+        val queued = collection.mutable.Queue.from(initialInvalidatedClassNames)
+
+        while (queued.nonEmpty) {
+          val current = queued.dequeue()
+          seenList.append(current)
+          seen.add(current)
+
+          for (next <- relations.usesInternalClass(current)) {
+            if (!seen.contains(next)) {
+              seen.add(next)
+              queued.enqueue(next)
+            }
+          }
+
+          for (next <- relations.usesExternal(current)) {
+            if (!seen.contains(next)) {
+              seen.add(next)
+              queued.enqueue(next)
+            }
+          }
+        }
+
+        seenList
+          .iterator
+          .flatMap { invalidatedClassName =>
+            relations.definesClass(invalidatedClassName)
+          }
+          .flatMap { source =>
+            relations.products(source)
+          }
+          .map { product =>
+            os.Path(product.id)
+          }
+          .toSet
+      }
+
+      val relations = combinedAnalysis.relations
+
+      val invalidatedAbsoluteClasspaths = getInvalidatedClasspaths(
+        result.changedSources.flatMap { source =>
+          relations.classNames(xsbti.VirtualFileRef.of(source))
+        },
+        combinedAnalysis.relations
+      )
+
+      // We only care about testing class, so we can:
+      // - filter out all class path that start with `testClasspath()`
+      // - strip the prefix and safely turn them into module class path
+      
+      val testClasspaths = testClasspath()
+      val invalidatedClassNames = invalidatedAbsoluteClasspaths.flatMap { absoluteClasspath =>
+        testClasspaths.collectFirst {
+          case path if absoluteClasspath.startsWith(path.path) =>
+            absoluteClasspath.relativeTo(path.path).segments.map(_.stripSuffix(".class")).mkString(".")
+        }
+      }
+      val testingClasses = invalidatedClassNames ++ failedTestClasses
+      val testClasses = testForkGrouping().map(_.filter(testingClasses.contains)).filter(_.nonEmpty)
+
+      // Clean up the directory for test runners
+      os.walk(Task.dest).foreach { subPath => os.remove.all(subPath) }
+
+      val quickTestReportXml = testReportXml()
+
+      val testModuleUtil = new TestModuleUtil(
+        testUseArgsFile(),
+        forkArgs(),
+        Seq.empty,
+        zincWorker().scalalibClasspath(),
+        resources(),
+        testFramework(),
+        runClasspath(),
+        testClasspaths,
+        args.toSeq,
+        testClasses,
+        zincWorker().testrunnerEntrypointClasspath(),
+        forkEnv(),
+        testSandboxWorkingDir(),
+        forkWorkingDir(),
+        quickTestReportXml,
+        zincWorker().javaHome().map(_.path),
+        testParallelism()
+      )
+
+      val results = testModuleUtil.runTests()
+
+      val badTestClasses = (results match {
+        case Result.Failure(_) =>
+          // Consider all quick testing classes as failed
+          testClasses.flatten
+        case Result.Success((_, results)) =>
+          // Get all test classes that failed
+          results
+            .filter(testResult => Set("Error", "Failure").contains(testResult.status))
+            .map(_.fullyQualifiedName)
+      }).distinct
+
+      os.write.over(transitiveStampsFile, upickle.default.write(result))
+      os.write.over(quicktestFailedClassesLog, upickle.default.write(badTestClasses))
+      os.write.over(invalidatedClassesLog, upickle.default.write(invalidatedClassNames))
+      results match {
+        case Result.Failure(errMsg) => Result.Failure(errMsg)
+        case Result.Success((doneMsg, results)) =>
+          try TestModule.handleResults(doneMsg, results, Task.ctx(), quickTestReportXml)
+          catch {
+            case e: Throwable => Result.Failure("Test reporting failed: " + e)
+          }
+      }
+    }
   }
 
   def defaultCommandName(): String = "run"

scalalib/src/mill/scalalib/TestModule.scala

@@ -205,10 +205,12 @@ trait TestModule
       globSelectors: Task[Seq[String]]
   ): Task[(String, Seq[TestResult])] =
     Task.Anon {
+      val testGlobSelectors = globSelectors()
+      val reportXml = testReportXml()
       val testModuleUtil = new TestModuleUtil(
         testUseArgsFile(),
         forkArgs(),
-        globSelectors(),
+        testGlobSelectors,
         jvmWorker().scalalibClasspath(),
         resources(),
         testFramework(),
@@ -220,12 +222,25 @@ trait TestModule
         forkEnv(),
         testSandboxWorkingDir(),
         forkWorkingDir(),
-        testReportXml(),
+        reportXml,
         jvmWorker().javaHome().map(_.path),
         testParallelism(),
         testLogLevel()
       )
-      testModuleUtil.runTests()
+      val result = testModuleUtil.runTests()
+
+      result match {
+        case Result.Failure(errMsg) => Result.Failure(errMsg)
+        case Result.Success((doneMsg, results)) =>
+          if (results.isEmpty && testGlobSelectors.nonEmpty) throw new Result.Exception(
+            s"Test selector does not match any test: ${testGlobSelectors.mkString(" ")}" +
+              "\nRun discoveredTestClasses to see available tests"
+          )
+          try TestModule.handleResults(doneMsg, results, Task.ctx(), reportXml)
+          catch {
+            case e: Throwable => Result.Failure("Test reporting failed: " + e)
+          }
+      }
     }
 
   /**

scalalib/src/mill/scalalib/TestModuleUtil.scala

@@ -102,21 +102,11 @@ private final class TestModuleUtil(
       }
     if (selectors.nonEmpty && filteredClassLists.isEmpty) throw doesNotMatchError
 
-    val result = if (testParallelism) {
+    if (testParallelism) {
       runTestQueueScheduler(filteredClassLists)
     } else {
       runTestDefault(filteredClassLists)
     }
-
-    result match {
-      case Result.Failure(errMsg) => Result.Failure(errMsg)
-      case Result.Success((doneMsg, results)) =>
-        if (results.isEmpty && selectors.nonEmpty) throw doesNotMatchError
-        try TestModuleUtil.handleResults(doneMsg, results, Task.ctx(), testReportXml)
-        catch {
-          case e: Throwable => Result.Failure("Test reporting failed: " + e)
-        }
-    }
   }
 
   private def callTestRunnerSubprocess(

testkit/src/mill/testkit/UnitTester.scala

@@ -52,7 +52,7 @@ class UnitTester(
     inStream: InputStream,
     debugEnabled: Boolean,
     env: Map[String, String],
-    resetSourcePath: Boolean
+    resetSourcePath: Boolean = true
 )(implicit fullName: sourcecode.FullName) extends AutoCloseable {
   val outPath: os.Path = module.moduleDir / "out"