Add DiGraph factory method and prettyTree

* New factory method enables direct construction of DiGraphs from edges
* DiGraph.prettyTree enables visualization of tree or multi-tree
  diagraphs

1 files changed, 47 insertions, 0 deletions

diff --git a/src/main/scala/firrtl/graph/DiGraph.scala b/src/main/scala/firrtl/graph/DiGraph.scala
index 3a08d05e..9f6ffeb2 100644
--- a/src/main/scala/firrtl/graph/DiGraph.scala
+++ b/src/main/scala/firrtl/graph/DiGraph.scala
@@ -4,6 +4,7 @@ package firrtl.graph
 
 import scala.collection.{mutable, Map, Set}
 import scala.collection.mutable.{LinkedHashMap, LinkedHashSet}
+import firrtl.options.DependencyManagerUtils.{CharSet, PrettyCharSet}
 
 /** An exception that is raised when an assumed DAG has a cycle */
 class CyclicException(val node: Any) extends Exception(s"No valid linearization for cyclic graph, found at $node")
@@ -31,6 +32,16 @@ object DiGraph {
     }
     new DiGraph(edgeDataCopy)
   }
+
+  /** Create a DiGraph from edges */
+  def apply[T](edges: (T, T)*): DiGraph[T] = {
+    val edgeMap = new LinkedHashMap[T, LinkedHashSet[T]]
+    for ((from, to) <- edges) {
+      val set = edgeMap.getOrElseUpdate(from, new LinkedHashSet[T])
+      set += to
+    }
+    new DiGraph(edgeMap)
+  }
 }
 
 /** Represents common behavior of all directed graphs */
@@ -386,6 +397,42 @@ class DiGraph[T](private[graph] val edges: LinkedHashMap[T, LinkedHashSet[T]]) {
     that.edges.foreach({ case (k, v) => eprime.getOrElseUpdate(k, new LinkedHashSet[T]) ++= v })
     new DiGraph(eprime)
   }
+
+  /** Serializes a `DiGraph[String]` as a pretty tree
+    *
+    * Multiple roots are supported, but cycles are not.
+    */
+  def prettyTree(charSet: CharSet = PrettyCharSet)(implicit ev: T =:= String): String = {
+    // Set up characters for building the tree
+    val (l, n, c) = (charSet.lastNode, charSet.notLastNode, charSet.continuation)
+    val ctab = " " * c.size + " "
+
+    // Recursively adds each node of the DiGraph to accumulating List[String]
+    // Uses List because prepend is cheap and this prevents quadratic behavior of String
+    //   concatenations or even flatMapping on Seqs
+    def rec(tab: String, node: T, mark: String, prev: List[String]): List[String] = {
+      val here = s"$mark$node"
+      val children = this.getEdges(node)
+      val last = children.size - 1
+      children.toList // Convert LinkedHashSet to List to avoid determinism issues
+        .zipWithIndex // Find last
+        .foldLeft(here :: prev) {
+          case (acc, (nodex, idx)) =>
+            val nextTab = if (idx == last) tab + ctab else tab + c + " "
+            val nextMark = if (idx == last) tab + l else tab + n
+            rec(nextTab, nodex, nextMark + " ", acc)
+        }
+    }
+    this.findSources
+      .toList // Convert LinkedHashSet to List to avoid determinism issues
+      .sortBy(_.toString) // Make order deterministic
+      .foldLeft(Nil: List[String]) {
+        case (acc, root) => rec("", root, "", acc)
+      }
+      .reverse
+      .mkString("\n")
+  }
+
 }
 
 class MutableDiGraph[T] extends DiGraph[T](new LinkedHashMap[T, LinkedHashSet[T]]) {