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

5 files changed, 144 insertions, 28 deletions

diff --git a/docs/VisualizePlan.svg b/docs/VisualizePlan.svg
new file mode 100644
index 00000000..d0e0a15b
--- /dev/null
+++ b/docs/VisualizePlan.svg
@@ -0,0 +1,49 @@
+<svg width="454pt" height="188pt"
+ viewBox="0.00 0.00 454.06 188.00" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
+<g id="graph0" class="graph" transform="scale(1 1) rotate(0) translate(4 184)">
+<title>example1</title>
+<polygon fill="#ffffff" stroke="transparent" points="-4,4 -4,-184 450.0647,-184 450.0647,4 -4,4"/>
+<!-- moduledefs.sources -->
+<g id="node1" class="node">
+<title>moduledefs.sources</title>
+<ellipse fill="none" stroke="#000000" stroke-dasharray="1,5" cx="87.4606" cy="-18" rx="87.4212" ry="18"/>
+<text text-anchor="middle" x="87.4606" y="-13.8" font-family="Times,serif" font-size="14.00" fill="#000000">moduledefs.sources</text>
+</g>
+<!-- moduledefs.generatedSources -->
+<g id="node2" class="node">
+<title>moduledefs.generatedSources</title>
+<ellipse fill="none" stroke="#000000" stroke-dasharray="1,5" cx="319.4606" cy="-18" rx="126.7082" ry="18"/>
+<text text-anchor="middle" x="319.4606" y="-13.8" font-family="Times,serif" font-size="14.00" fill="#000000">moduledefs.generatedSources</text>
+</g>
+<!-- moduledefs.allSources -->
+<g id="node3" class="node">
+<title>moduledefs.allSources</title>
+<ellipse fill="none" stroke="#000000" stroke-dasharray="1,5" cx="203.4606" cy="-90" rx="98.9552" ry="18"/>
+<text text-anchor="middle" x="203.4606" y="-85.8" font-family="Times,serif" font-size="14.00" fill="#000000">moduledefs.allSources</text>
+</g>
+<!-- moduledefs.allSources&#45;&gt;moduledefs.sources -->
+<g id="edge1" class="edge">
+<title>moduledefs.allSources&#45;&gt;moduledefs.sources</title>
+<path fill="none" stroke="#000000" d="M175.3802,-72.5708C159.9918,-63.0194 140.734,-51.0662 124.1675,-40.7836"/>
+<polygon fill="#000000" stroke="#000000" points="125.6246,-37.5686 115.2824,-35.2687 121.933,-43.5161 125.6246,-37.5686"/>
+</g>
+<!-- moduledefs.allSources&#45;&gt;moduledefs.generatedSources -->
+<g id="edge2" class="edge">
+<title>moduledefs.allSources&#45;&gt;moduledefs.generatedSources</title>
+<path fill="none" stroke="#000000" d="M231.541,-72.5708C246.7031,-63.1599 265.6216,-51.4173 282.0211,-41.2383"/>
+<polygon fill="#000000" stroke="#000000" points="284.1823,-44.0163 290.833,-35.7689 280.4908,-38.0688 284.1823,-44.0163"/>
+</g>
+<!-- moduledefs.allSourceFiles -->
+<g id="node4" class="node">
+<title>moduledefs.allSourceFiles</title>
+<ellipse fill="none" stroke="#000000" cx="203.4606" cy="-162" rx="114" ry="18"/>
+<text text-anchor="middle" x="203.4606" y="-157.8" font-family="Times,serif" font-size="14.00" fill="#000000">moduledefs.allSourceFiles</text>
+</g>
+<!-- moduledefs.allSourceFiles&#45;&gt;moduledefs.allSources -->
+<g id="edge3" class="edge">
+<title>moduledefs.allSourceFiles&#45;&gt;moduledefs.allSources</title>
+<path fill="none" stroke="#000000" d="M203.4606,-143.8314C203.4606,-136.131 203.4606,-126.9743 203.4606,-118.4166"/>
+<polygon fill="#000000" stroke="#000000" points="206.9607,-118.4132 203.4606,-108.4133 199.9607,-118.4133 206.9607,-118.4132"/>
+</g>
+</g>
+</svg>
diff --git a/docs/pages/1 - Intro to Mill.md b/docs/pages/1 - Intro to Mill.md
index a3300595..49ac64da 100644
--- a/docs/pages/1 - Intro to Mill.md
+++ b/docs/pages/1 - Intro to Mill.md
@@ -558,6 +558,29 @@ compilation output:
 
 ![VisualizeCompile.svg](VisualizeCompile.svg)
 
+### visualizePlan
+
+```bash
+$ mill show visualizePlan moduledefs.allSourceFiles
+[
+    ".../out/visualizePlan/dest/out.txt",
+    ".../out/visualizePlan/dest/out.dot",
+    ".../out/visualizePlan/dest/out.json",
+    ".../out/visualizePlan/dest/out.png",
+    ".../out/visualizePlan/dest/out.svg"
+]
+```
+
+`mill show visualizePlan` is similar except that it shows a graph of the entire
+build plan, including dependencies not directly resolved by the query. Targets
+directly resolved are shown with a solid border, and dependencies are shown with
+a dotted border.
+
+The above command generates the following diagram:
+
+![VisualizePlan.svg](VisualizePlan.svg)
+
+Currently, visualizePlan does not support graphs that contain external modules.
 
 ### clean
 
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)