Visualize Plan (#404)

* Make necessary import changes

* Refactor to allow calling internally w/o println

* Refactor to allow multiple visualize modes

* Add new visualizaPlan grap entire plan

* Remove and alphabetize imports

* Document visualizePlan

3 files changed, 72 insertions, 28 deletions

diff --git a/main/graphviz/src/mill/main/graphviz/GraphvizTools.scala b/main/graphviz/src/mill/main/graphviz/GraphvizTools.scala
index 300bd98a..4e8c59ce 100644
--- a/main/graphviz/src/mill/main/graphviz/GraphvizTools.scala
+++ b/main/graphviz/src/mill/main/graphviz/GraphvizTools.scala
@@ -1,17 +1,18 @@
 package mill.main.graphviz
 import ammonite.ops.Path
+import guru.nidi.graphviz.attribute.Style
 import mill.define.{Graph, NamedTask}
 import org.jgrapht.graph.{DefaultEdge, SimpleDirectedGraph}
 object GraphvizTools{
-  def apply(rs: Seq[NamedTask[Any]], dest: Path) = {
+  def apply(targets: Seq[NamedTask[Any]], rs: Seq[NamedTask[Any]], dest: Path) = {
     val transitive = Graph.transitiveTargets(rs.distinct)
     val topoSorted = Graph.topoSorted(transitive)
     val goalSet = rs.toSet
     val sortedGroups = Graph.groupAroundImportantTargets(topoSorted){
       case x: NamedTask[Any] if goalSet.contains(x) => x
     }
-    import guru.nidi.graphviz.model.Factory._
     import guru.nidi.graphviz.engine.{Format, Graphviz}
+    import guru.nidi.graphviz.model.Factory._
 
     val edgesIterator =
       for((k, vs) <- sortedGroups.items())
@@ -38,7 +39,12 @@ object GraphvizTools{
 
 
     org.jgrapht.alg.TransitiveReduction.INSTANCE.reduce(jgraph)
-    val nodes = indexToTask.map(t => node(t.ctx.segments.render))
+    val nodes = indexToTask.map(t =>
+      node(t.ctx.segments.render).`with`{
+        if(targets.contains(t)) Style.SOLID
+        else Style.DOTTED
+      }
+    )
 
     var g = graph("example1").directed
     for(i <- indexToTask.indices){
diff --git a/main/src/mill/main/MainModule.scala b/main/src/mill/main/MainModule.scala
index f7d4e7b9..fbf89a6d 100644
--- a/main/src/mill/main/MainModule.scala
+++ b/main/src/mill/main/MainModule.scala
@@ -1,15 +1,14 @@
 package mill.main
 
+import java.util.concurrent.LinkedBlockingQueue
+
 import ammonite.ops.Path
-import coursier.Cache
-import coursier.maven.MavenRepository
 import mill.T
-import mill.define.{Graph, NamedTask, Task}
+import mill.define.{NamedTask, Task}
 import mill.eval.{Evaluator, PathRef, Result}
-import mill.util.{Loose, PrintLogger, Watched}
+import mill.util.{Ctx, PrintLogger, Watched}
 import pprint.{Renderer, Truncated}
 import upickle.Js
-import mill.util.JsonFormatters._
 object MainModule{
   def resolveTasks[T](evaluator: Evaluator[Any], targets: Seq[String], multiSelect: Boolean)
                      (f: List[NamedTask[Any]] => T) = {
@@ -66,21 +65,29 @@ trait MainModule extends mill.Module{
     * executed in what order, without actually executing them.
     */
   def plan(evaluator: Evaluator[Any], targets: String*) = mill.T.command{
+    plan0(evaluator, targets) match{
+      case Right(success) => {
+        success.foreach(println)
+        Result.Success(success)
+      }
+      case Left(err) => Result.Failure(err)
+    }
+  }
+
+  private def plan0(evaluator: Evaluator[Any], targets: Seq[String]) = {
     val resolved = RunScript.resolveTasks(
       mill.main.ResolveTasks, evaluator, targets, multiSelect = true
     )
 
-    resolved match{
-      case Left(err) => Result.Failure(err)
+    resolved match {
+      case Left(err) => Left(err)
       case Right(rs) =>
         val (sortedGroups, transitive) = Evaluator.plan(evaluator.rootModule, rs)
-        val labels = sortedGroups
+        Right(sortedGroups
           .keys()
           .collect{ case Right(r) => r.segments.render}
           .toArray
-
-        labels.foreach(println)
-        Result.Success(labels)
+        )
     }
   }
 
@@ -221,15 +228,46 @@ trait MainModule extends mill.Module{
   }
 
   def visualize(evaluator: Evaluator[Any], targets: String*) = mill.T.command{
-    val resolved = RunScript.resolveTasks(
-      mill.main.ResolveTasks, evaluator, targets, multiSelect = true
-    )
-    resolved match{
+    visualize0(evaluator, targets, T.ctx(), mill.main.VisualizeModule.worker())
+  }
+
+  def visualizePlan(evaluator: Evaluator[Any], targets: String*) = mill.T.command{
+    plan0(evaluator, targets) match {
+      case Right(planResults) =>
+        visualize0(evaluator, targets, T.ctx(), mill.main.VisualizeModule.worker(), Some(planResults))
       case Left(err) => Result.Failure(err)
-      case Right(rs) =>
-        val (in, out) = mill.main.VisualizeModule.worker()
-        in.put((rs, T.ctx().dest))
-        out.take()
     }
   }
+
+  private type VizWorker = (LinkedBlockingQueue[(scala.Seq[_], scala.Seq[_], Path)],
+    LinkedBlockingQueue[Result[scala.Seq[PathRef]]])
+
+  private def visualize0(evaluator: Evaluator[Any], targets: Seq[String], ctx: Ctx, vizWorker: VizWorker,
+                         planTasks: Option[Array[String]] = None) = {
+    def resolveTasks(targets: Seq[String]): Either[String, List[NamedTask[Any]]] = {
+      RunScript.resolveTasks(
+        mill.main.ResolveTasks, evaluator, targets, multiSelect = true
+      )
+    }
+
+    def callVisualizeModule(rs: List[NamedTask[Any]], allRs: List[NamedTask[Any]]) = {
+      val (in, out) = vizWorker
+      in.put((rs, allRs, ctx.dest))
+      out.take()
+    }
+
+    val resolved = resolveTasks(targets)
+
+    resolved match {
+      case Left(err) => Result.Failure(err)
+      case Right(rs) => planTasks match {
+        case Some(allTasks) => resolveTasks(allTasks) match {
+          case Left (err) => Result.Failure (err)
+          case Right (allRs) => callVisualizeModule (rs, allRs)
+        }
+        case None => callVisualizeModule(rs, rs)
+      }
+    }
+  }
+
 }
diff --git a/main/src/mill/main/VisualizeModule.scala b/main/src/mill/main/VisualizeModule.scala
index 8e022a33..6945cf2f 100644
--- a/main/src/mill/main/VisualizeModule.scala
+++ b/main/src/mill/main/VisualizeModule.scala
@@ -6,9 +6,9 @@ import ammonite.ops.Path
 import coursier.Cache
 import coursier.core.Repository
 import coursier.maven.MavenRepository
-import mill.{T, eval}
+import mill.T
 import mill.define.{Discover, ExternalModule}
-import mill.eval.{Evaluator, PathRef, Result}
+import mill.eval.{PathRef, Result}
 
 
 object VisualizeModule extends ExternalModule with VisualizeModule {
@@ -36,7 +36,7 @@ trait VisualizeModule extends mill.define.TaskModule{
     * can communicate via in/out queues.
     */
   def worker = T.worker{
-    val in = new LinkedBlockingQueue[(Seq[_], Path)]()
+    val in = new LinkedBlockingQueue[(Seq[_], Seq[_], Path)]()
     val out = new LinkedBlockingQueue[Result[Seq[PathRef]]]()
 
     val cl = mill.util.ClassLoader.create(
@@ -46,10 +46,10 @@ trait VisualizeModule extends mill.define.TaskModule{
     val visualizeThread = new java.lang.Thread(() =>
       while(true){
         val res = Result.create{
-          val (tasks, dest) = in.take()
+          val (targets, tasks, dest) = in.take()
           cl.loadClass("mill.main.graphviz.GraphvizTools")
-            .getMethod("apply", classOf[Seq[_]], classOf[Path])
-            .invoke(null, tasks, dest)
+            .getMethod("apply", classOf[Seq[_]], classOf[Seq[_]], classOf[Path])
+            .invoke(null, targets, tasks, dest)
             .asInstanceOf[Seq[PathRef]]
         }
         out.put(res)