Tests are passing under the new group evaluation model

2 files changed, 46 insertions, 83 deletions

diff --git a/core/src/main/scala/mill/eval/Evaluator.scala b/core/src/main/scala/mill/eval/Evaluator.scala
index da6e2656..8163a304 100644
--- a/core/src/main/scala/mill/eval/Evaluator.scala
+++ b/core/src/main/scala/mill/eval/Evaluator.scala
@@ -10,44 +10,43 @@ import scala.collection.mutable
 class Evaluator(workspacePath: Path,
                 labeling: Map[Task[_], Labelled[_]]){
 
-  def evaluate(targets: OSet[Task[_]]): Evaluator.Results = {
+  def evaluate(goals: OSet[Task[_]]): Evaluator.Results = {
     mkdir(workspacePath)
 
-    val transitive = Evaluator.transitiveTargets(targets)
+    val transitive = Evaluator.transitiveTargets(goals)
     val topoSorted = Evaluator.topoSorted(transitive)
-    val sortedGroups = Evaluator.groupAroundNamedTargets(topoSorted, labeling)
+    val sortedGroups = Evaluator.groupAroundImportantTargets(
+      topoSorted,
+      t => labeling.contains(t) || goals.contains(t)
+    )
 
     val evaluated = new OSet.Mutable[Task[_]]
     val results = mutable.LinkedHashMap.empty[Task[_], Any]
 
-    for (groupIndex <- sortedGroups.keys()){
-      val group = sortedGroups.lookupKey(groupIndex)
-
-      val (newResults, newEvaluated) = evaluateGroupCached(
-        group,
-        results,
-        sortedGroups
-      )
-      evaluated.appendAll(newEvaluated)
+    for ((terminal, group)<- sortedGroups.items()){
+      val (newResults, newEvaluated) = evaluateGroupCached(terminal, group, results)
+      for(ev <- newEvaluated){
+        evaluated.append(ev)
+      }
       for((k, v) <- newResults) results.put(k, v)
 
     }
 
-    Evaluator.Results(targets.items.map(results), evaluated, transitive)
+    Evaluator.Results(goals.items.map(results), evaluated, transitive)
   }
 
-  def evaluateGroupCached(group: OSet[Task[_]],
-                          results: collection.Map[Task[_], Any],
-                          sortedGroups: MultiBiMap[Int, Task[_]]): (collection.Map[Task[_], Any], Seq[Task[_]]) = {
+  def evaluateGroupCached(terminal: Task[_],
+                          group: OSet[Task[_]],
+                          results: collection.Map[Task[_], Any]): (collection.Map[Task[_], Any], Seq[Task[_]]) = {
 
 
-    val (externalInputs, terminals) = partitionGroupInputOutput(group, results)
+    val externalInputs = group.items.flatMap(_.inputs).filter(!group.contains(_))
 
     val inputsHash =
       externalInputs.toIterator.map(results).toVector.hashCode +
       group.toIterator.map(_.sideHash).toVector.hashCode()
 
-    val (targetDestPath, metadataPath) = labeling.get(terminals.items(0)) match{
+    val (targetDestPath, metadataPath) = labeling.get(terminal) match{
       case Some(labeling) =>
         val targetDestPath = workspacePath / labeling.segments
         val metadataPath = targetDestPath / up / (targetDestPath.last + ".mill.json")
@@ -58,16 +57,14 @@ class Evaluator(workspacePath: Path,
     val cached = for{
       metadataPath <- metadataPath
       json <- scala.util.Try(Json.parse(read.getInputStream(metadataPath))).toOption
-      (cachedHash, terminalResults) <- Json.fromJson[(Int, Seq[JsValue])](json).asOpt
+      (cachedHash, terminalResult) <- Json.fromJson[(Int, JsValue)](json).asOpt
       if cachedHash == inputsHash
-    } yield terminalResults
+    } yield terminalResult
 
     cached match{
-      case Some(terminalResults) =>
+      case Some(terminalResult) =>
         val newResults = mutable.LinkedHashMap.empty[Task[_], Any]
-        for((terminal, res) <- terminals.items.zip(terminalResults)){
-          newResults(terminal) = labeling(terminal).format.reads(res).get
-        }
+        newResults(terminal) = labeling(terminal).format.reads(terminalResult).get
         (newResults, Nil)
 
       case _ =>
@@ -76,13 +73,13 @@ class Evaluator(workspacePath: Path,
         if (labeled.nonEmpty){
           println(fansi.Color.Blue("Running " + labeled.map(_.segments.mkString(".")).mkString(", ")))
         }
-        val (newResults, newEvaluated, terminalResults) = evaluateGroup(group, results, targetDestPath)
+        val (newResults, newEvaluated, terminalResult) = evaluateGroup(group, results, targetDestPath)
 
         metadataPath.foreach(
           write.over(
             _,
             Json.prettyPrint(
-              Json.toJson(inputsHash -> terminals.toList.map(terminalResults))
+              Json.toJson(inputsHash -> terminalResult)
             ),
           )
         )
@@ -91,24 +88,17 @@ class Evaluator(workspacePath: Path,
     }
   }
 
-  def partitionGroupInputOutput(group: OSet[Task[_]],
-                                results: collection.Map[Task[_], Any]) = {
-    val allInputs = group.items.flatMap(_.inputs)
-    val (internalInputs, externalInputs) = allInputs.partition(group.contains)
-    val internalInputSet = internalInputs.toSet
-    val terminals = group.filter(x => !internalInputSet(x) || labeling.contains(x))
-    (OSet.from(externalInputs.distinct), terminals)
-  }
 
   def evaluateGroup(group: OSet[Task[_]],
                     results: collection.Map[Task[_], Any],
                     targetDestPath: Option[Path]) = {
 
     targetDestPath.foreach(rm)
-    val terminalResults = mutable.LinkedHashMap.empty[Task[_], JsValue]
+    var terminalResult: JsValue = null
     val newEvaluated = mutable.Buffer.empty[Task[_]]
     val newResults = mutable.LinkedHashMap.empty[Task[_], Any]
-    for (target <- group.items) {
+    for (target <- group.items if !results.contains(target)) {
+
       newEvaluated.append(target)
       val targetInputValues = target.inputs.toVector.map(x =>
         newResults.getOrElse(x, results(x))
@@ -117,7 +107,7 @@ class Evaluator(workspacePath: Path,
       val args = new Args(targetInputValues, targetDestPath.orNull)
       val res = target.evaluate(args)
       for(targetLabel <- labeling.get(target)){
-        terminalResults(target) = targetLabel
+        terminalResult = targetLabel
           .format
           .asInstanceOf[Format[Any]]
           .writes(res.asInstanceOf[Any])
@@ -125,7 +115,7 @@ class Evaluator(workspacePath: Path,
       newResults(target) = res
     }
 
-    (newResults, newEvaluated, terminalResults)
+    (newResults, newEvaluated, terminalResult)
   }
 
 }
@@ -134,52 +124,23 @@ class Evaluator(workspacePath: Path,
 object Evaluator{
   class TopoSorted private[Evaluator](val values: OSet[Task[_]])
   case class Results(values: Seq[Any], evaluated: OSet[Task[_]], transitive: OSet[Task[_]])
-  def groupAroundNamedTargets(topoSortedTargets: TopoSorted,
-                              labeling: Map[Task[_], Labelled[_]]): MultiBiMap[Int, Task[_]] = {
-
-    val grouping = new MultiBiMap.Mutable[Int, Task[_]]()
-
-    var groupCount = 0
-
-    for(target <- topoSortedTargets.values.items.reverseIterator){
-      if (!grouping.containsValue(target)){
-        grouping.add(groupCount, target)
-        groupCount += 1
-      }
-
-      val targetGroup = grouping.lookupValue(target)
-      for(upstream <- target.inputs){
-        grouping.lookupValueOpt(upstream) match{
-          case None if !labeling.contains(upstream) =>
-            grouping.add(targetGroup, upstream)
-          case Some(upstreamGroup)
-            if !labeling.contains(upstream) && upstreamGroup != targetGroup =>
-            val upstreamTargets = grouping.removeAll(upstreamGroup)
-            grouping.addAll(targetGroup, upstreamTargets)
-          case _ => //donothing
+  def groupAroundImportantTargets(topoSortedTargets: TopoSorted,
+                                  important: Task[_] => Boolean): MultiBiMap[Task[_], Task[_]] = {
+
+    val output = new MultiBiMap.Mutable[Task[_], Task[_]]()
+    for (target <- topoSortedTargets.values if important(target)) {
+
+      val transitiveTargets = new OSet.Mutable[Task[_]]
+      def rec(t: Task[_]): Unit = {
+        if (transitiveTargets.contains(t)) () // do nothing
+        else if (important(t) && t != target) () // do nothing
+        else {
+          transitiveTargets.append(t)
+          t.inputs.foreach(rec)
         }
       }
-    }
-
-
-    // Sort groups amongst themselves, and sort the contents of each group
-    // before aggregating it into the final output
-    val groupGraph = mutable.Buffer.fill[Seq[Int]](groupCount)(Nil)
-    for((groupId, groupTasks) <- grouping.items()){
-      groupGraph(groupId) =
-        groupTasks.toIterator.flatMap(_.inputs).map(grouping.lookupValue).toArray.distinct.toSeq
-    }
-    // Given input topoSortedTargets has no cycles, group graph should not have cycles
-    val groupOrdering = Tarjans.apply(groupGraph)
-
-
-    val targetOrdering = topoSortedTargets.values.items.zipWithIndex.toMap
-    val output = new MultiBiMap.Mutable[Int, Task[_]]
-    for((groupIndices, i) <- groupOrdering.zipWithIndex){
-      val sortedGroup = OSet.from(
-        groupIndices.flatMap(grouping.lookupKeyOpt).flatten.sortBy(targetOrdering)
-      )
-      output.addAll(i, sortedGroup)
+      rec(target)
+      output.addAll(target, topoSorted(transitiveTargets).values)
     }
     output
   }
@@ -207,7 +168,7 @@ object Evaluator{
 
     val numberedEdges =
       for(t <- transitiveTargets.items)
-      yield t.inputs.map(targetIndices)
+      yield t.inputs.collect(targetIndices)
 
     val sortedClusters = Tarjans(numberedEdges)
     val nonTrivialClusters = sortedClusters.filter(_.length > 1)
diff --git a/core/src/main/scala/mill/util/MultiBiMap.scala b/core/src/main/scala/mill/util/MultiBiMap.scala
index 2052ef90..ef5359bc 100644
--- a/core/src/main/scala/mill/util/MultiBiMap.scala
+++ b/core/src/main/scala/mill/util/MultiBiMap.scala
@@ -16,9 +16,11 @@ trait MultiBiMap[K, V]{
   def removeAll(k: K): OSet[V]
   def addAll(k: K, vs: TraversableOnce[V]): Unit
   def keys(): Iterator[K]
+  def items(): Iterator[(K, OSet[V])]
   def values(): Iterator[OSet[V]]
   def keyCount: Int
 }
+
 object MultiBiMap{
   class Mutable[K, V]() extends MultiBiMap[K, V]{
     private[this] val valueToKey = mutable.LinkedHashMap.empty[V, K]