Add selective.resolveTree task to help debug selective execution misbehaviors (#4349)

Also lots of general refactoring for the
SelectiveExecution/SpanningForest logic

I need this myself to help debug
#4343

Author: lihaoyi

build.mill

@@ -609,7 +609,8 @@ trait MillStableScalaModule extends MillPublishScalaModule with Mima {
     ProblemFilter.exclude[Problem]("*.bspJvmBuildTarget"),
     ProblemFilter.exclude[Problem]("mill.scalalib.RunModule#RunnerImpl.*"),
     ProblemFilter.exclude[Problem]("mill.util.PromptLogger#*"),
-    ProblemFilter.exclude[Problem]("mill.util.PromptLoggerUtil.*")
+    ProblemFilter.exclude[Problem]("mill.util.PromptLoggerUtil.*"),
+    ProblemFilter.exclude[Problem]("mill.main.SelectiveExecution*")
   )
   def mimaPreviousVersions: T[Seq[String]] = Settings.mimaBaseVersions
 

example/large/selective/9-selective-execution/build.mill

@@ -133,6 +133,50 @@ Test run bar.BarTests finished: 0 failed, 0 ignored, 1 total, ...
 //   "foo.test" [color=red penwidth=2]
 // }
 // ```
+//
+// For a more detailed report of how the changed inputs resulted in the selected tasks
+// being chosen, you can also use `selective.resolveTree` to print out the selected tasks
+// as a JSON tree illustrating the relationships between the invalidated inputs (at the root
+// of the tree) and the selected tasks (at the leaves of the tree)
+//
+
+/** Usage
+
+> mill selective.resolveTree __.test
+{
+  "bar.sources": {
+    "bar.test.sources": {
+      "bar.test.allSources": {
+        "bar.test.allSourceFiles": {
+          "bar.test.compile": {
+            "bar.test.localRunClasspath": {
+              "bar.test.testClasspath": {
+                "bar.test.test": {}
+              }
+            }
+          }
+        }
+      }
+    },
+    "bar.allSources": {
+      "bar.allSourceFiles": {
+        "bar.compile": {
+          "bar.localRunClasspath": {
+            "bar.localClasspath": {
+              "foo.test.transitiveLocalClasspath": {
+                "foo.test.runClasspath": {
+                  "foo.test.test": {}
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+*/
 
 // Similarly, if we make a change `qux/`, using selective execution will only run tests
 // in `qux.test`, and skip those in `foo.test` and `bar.test`.

main/codesig/src/ResolvedCalls.scala

@@ -1,6 +1,7 @@
 package mill.codesig
 import JvmModel._
 import JType.{Cls => JCls}
+import mill.util.SpanningForest
 import mill.util.SpanningForest.breadthFirst
 import upickle.default.{ReadWriter, macroRW}
 
@@ -31,12 +32,7 @@ object ResolvedCalls {
     val allDirectAncestors =
       localSummary.mapValues(_.directAncestors) ++ externalSummary.directAncestors
 
-    val directDescendents = {
-      allDirectAncestors
-        .toVector
-        .flatMap { case (k, vs) => vs.map((_, k)) }
-        .groupMap(_._1)(_._2)
-    }
+    val directDescendents = SpanningForest.reverseEdges(allDirectAncestors)
 
     // Given an external class, what are the local classes that inherit from it,
     // and what local methods may end up being called by the external class code
@@ -122,16 +118,13 @@ object ResolvedCalls {
 
       val allSamImplementors0 = allSamDefiners
         .toSeq
-        .flatMap { case (cls, sig) =>
-          breadthFirst(Seq(cls))(cls => directDescendents.getOrElse(cls, Nil)).map(_ -> sig)
+        .map { case (cls, sig) =>
+          sig -> breadthFirst(Seq(cls))(cls => directDescendents.getOrElse(cls, Nil))
         }
 
-      val allSamImplementors = allSamImplementors0
-        .groupMap(_._1)(_._2)
-        .view.mapValues(_.toSet)
-        .toMap
+      val allSamImplementors = mill.util.SpanningForest.reverseEdges(allSamImplementors0)
 
-      allSamImplementors
+      allSamImplementors.mapValues(_.toSet).toMap
     }
 
     val localCalls = {

main/eval/src/mill/eval/EvaluatorCore.scala

@@ -72,14 +72,13 @@ private[mill] trait EvaluatorCore extends GroupEvaluator {
 
     val threadNumberer = new ThreadNumberer()
     val (sortedGroups, transitive) = Plan.plan(goals)
-    val interGroupDeps = findInterGroupDeps(sortedGroups)
+    val interGroupDeps = EvaluatorCore.findInterGroupDeps(sortedGroups)
     val terminals0 = sortedGroups.keys().toVector
     val failed = new AtomicBoolean(false)
     val count = new AtomicInteger(1)
     val indexToTerminal = sortedGroups.keys().toArray
-    val terminalToIndex = indexToTerminal.zipWithIndex.toMap
 
-    EvaluatorLogs.logDependencyTree(interGroupDeps, indexToTerminal, terminalToIndex, outPath)
+    EvaluatorLogs.logDependencyTree(interGroupDeps, indexToTerminal, outPath)
 
     // Prepare a lookup tables up front of all the method names that each class owns,
     // and the class hierarchy, so during evaluation it is cheap to look up what class
@@ -231,7 +230,6 @@ private[mill] trait EvaluatorCore extends GroupEvaluator {
     EvaluatorLogs.logInvalidationTree(
       interGroupDeps,
       indexToTerminal,
-      terminalToIndex,
       outPath,
       uncached,
       changedValueHash
@@ -262,8 +260,10 @@ private[mill] trait EvaluatorCore extends GroupEvaluator {
       results.map { case (k, v) => (k, v.map(_._1)) }
     )
   }
+}
 
-  private def findInterGroupDeps(sortedGroups: MultiBiMap[Terminal, Task[_]])
+private[mill] object EvaluatorCore {
+  def findInterGroupDeps(sortedGroups: MultiBiMap[Terminal, Task[_]])
       : Map[Terminal, Seq[Terminal]] = {
     sortedGroups
       .items()
@@ -277,10 +277,6 @@ private[mill] trait EvaluatorCore extends GroupEvaluator {
       }
       .toMap
   }
-}
-
-private[mill] object EvaluatorCore {
-
   case class Results(
       rawValues: Seq[Result[Val]],
       evaluated: Agg[Task[_]],

main/eval/src/mill/eval/EvaluatorLogs.scala

@@ -10,38 +10,33 @@ private[mill] object EvaluatorLogs {
   def logDependencyTree(
       interGroupDeps: Map[Terminal, Seq[Terminal]],
       indexToTerminal: Array[Terminal],
-      terminalToIndex: Map[Terminal, Int],
       outPath: os.Path
   ): Unit = {
+    val (vertexToIndex, edgeIndices) =
+      SpanningForest.graphMapToIndices(indexToTerminal, interGroupDeps)
+
     SpanningForest.writeJsonFile(
       outPath / OutFiles.millDependencyTree,
-      indexToTerminal.map(t => interGroupDeps.getOrElse(t, Nil).map(terminalToIndex).toArray),
+      edgeIndices,
       indexToTerminal.indices.toSet,
       indexToTerminal(_).render
     )
   }
   def logInvalidationTree(
       interGroupDeps: Map[Terminal, Seq[Terminal]],
       indexToTerminal: Array[Terminal],
-      terminalToIndex: Map[Terminal, Int],
       outPath: os.Path,
       uncached: ConcurrentHashMap[Terminal, Unit],
       changedValueHash: ConcurrentHashMap[Terminal, Unit]
   ): Unit = {
-
-    val reverseInterGroupDeps = interGroupDeps
-      .iterator
-      .flatMap { case (k, vs) => vs.map(_ -> k) }
-      .toSeq
-      .groupMap(_._1)(_._2)
+    val reverseInterGroupDeps = SpanningForest.reverseEdges(interGroupDeps)
 
     val changedTerminalIndices = changedValueHash.keys().asScala.toSet
-    val downstreamIndexEdges = indexToTerminal
-      .map(t =>
-        if (changedTerminalIndices(t))
-          reverseInterGroupDeps.getOrElse(t, Nil).map(terminalToIndex).toArray
-        else Array.empty[Int]
-      )
+
+    val (vertexToIndex, downstreamIndexEdges) = SpanningForest.graphMapToIndices(
+      indexToTerminal,
+      reverseInterGroupDeps.filterKeys(changedTerminalIndices).toMap
+    )
 
     val edgeSourceIndices = downstreamIndexEdges
       .zipWithIndex
@@ -53,7 +48,7 @@ private[mill] object EvaluatorLogs {
       downstreamIndexEdges,
       uncached.keys().asScala
         .flatMap { uncachedTask =>
-          val uncachedIndex = terminalToIndex(uncachedTask)
+          val uncachedIndex = vertexToIndex(uncachedTask)
           Option.when(
             // Filter out input and source tasks which do not cause downstream invalidations
             // from the invalidation tree, because most of them are un-interesting and the

main/src/mill/main/MainModule.scala

@@ -2,7 +2,7 @@ package mill.main
 
 import mill.api._
 import mill.define._
-import mill.eval.{Evaluator, EvaluatorPaths}
+import mill.eval.{Evaluator, EvaluatorPaths, Terminal}
 import mill.moduledefs.Scaladoc
 import mill.resolve.SelectMode.Separated
 import mill.resolve.{Resolve, SelectMode}
@@ -76,6 +76,23 @@ object MainModule {
         Result.Success(output)
     }
   }
+
+  def plan0(
+      evaluator: Evaluator,
+      tasks: Seq[String]
+  ): Either[String, Array[Terminal.Labelled[_]]] = {
+    Resolve.Tasks.resolve(
+      evaluator.rootModule,
+      tasks,
+      SelectMode.Multi
+    ) match {
+      case Left(err) => Left(err)
+      case Right(rs) =>
+        val (sortedGroups, _) = evaluator.plan(rs)
+        Right(sortedGroups.keys().collect { case r: Terminal.Labelled[_] => r }.toArray)
+    }
+  }
+
 }
 
 /**
@@ -121,7 +138,7 @@ trait MainModule extends BaseModule0 {
    */
   def plan(evaluator: Evaluator, targets: String*): Command[Array[String]] =
     Task.Command(exclusive = true) {
-      SelectiveExecution.plan0(evaluator, targets) match {
+      MainModule.plan0(evaluator, targets) match {
         case Left(err) => Result.Failure(err)
         case Right(success) =>
           val renderedTasks = success.map(_.segments.render)
@@ -516,7 +533,7 @@ trait MainModule extends BaseModule0 {
    */
   def visualizePlan(evaluator: Evaluator, targets: String*): Command[Seq[PathRef]] =
     Task.Command(exclusive = true) {
-      SelectiveExecution.plan0(evaluator, targets) match {
+      MainModule.plan0(evaluator, targets) match {
         case Left(err) => Result.Failure(err)
         case Right(planResults) => visualize0(
             evaluator,

main/src/mill/main/RunScript.scala

@@ -72,21 +72,16 @@ object RunScript {
       if (
         selectiveExecutionEnabled && os.exists(evaluator.outPath / OutFiles.millSelectiveExecution)
       ) {
-        SelectiveExecution
-          .diffMetadata(evaluator, targets.map(terminals(_).render).toSeq)
-          .map { x =>
-            val (selected, results) = x
-            val selectedSet = selected.toSet
-            (
-              targets.filter(t => t.isExclusiveCommand || selectedSet(terminals(t).render)),
-              results
-            )
-          }
-      } else Right(targets -> Map.empty)
+        val changedTasks = SelectiveExecution.computeChangedTasks0(evaluator, targets.toSeq)
+        val selectedSet = changedTasks.downstreamTasks.map(_.ctx.segments.render).toSet
+        (
+          targets.filter(t => t.isExclusiveCommand || selectedSet(terminals(t).render)),
+          changedTasks.results
+        )
+      } else (targets -> Map.empty)
 
     selectedTargetsOrErr match {
-      case Left(err) => (Nil, Left(err))
-      case Right((selectedTargets, selectiveResults)) =>
+      case (selectedTargets, selectiveResults) =>
         val evaluated: Results = evaluator.evaluate(selectedTargets, serialCommandExec = true)
         val watched = (evaluated.results.iterator ++ selectiveResults)
           .collect {