1 files changed, 42 insertions, 16 deletions

diff --git a/src/test/scala/firrtlTests/graph/DiGraphTests.scala b/src/test/scala/firrtlTests/graph/DiGraphTests.scala
index a0f45c80..52ded253 100644
--- a/src/test/scala/firrtlTests/graph/DiGraphTests.scala
+++ b/src/test/scala/firrtlTests/graph/DiGraphTests.scala
@@ -40,31 +40,45 @@ class DiGraphTests extends FirrtlFlatSpec {
 
   val degenerateGraph = DiGraph(Map("a" -> Set.empty[String]))
 
-  acyclicGraph.findSCCs.filter(_.length > 1) shouldBe empty
-
-  cyclicGraph.findSCCs.filter(_.length > 1) should not be empty
+  "A graph without cycles" should "have NOT SCCs" in {
+    acyclicGraph.findSCCs.filter(_.length > 1) shouldBe empty
+  }
 
-  acyclicGraph.path("a","e") should not be empty
+  "A graph with cycles" should "have SCCs" in {
+    cyclicGraph.findSCCs.filter(_.length > 1) should not be empty
+  }
 
-  an [PathNotFoundException] should be thrownBy acyclicGraph.path("e","a")
+  "Asking a DiGraph for a path that exists" should "work" in {
+    acyclicGraph.path("a","e") should not be empty
+  }
 
-  acyclicGraph.linearize.head should equal ("a")
+  "Asking a DiGraph for a path from one node to another with no path" should "error" in {
+    an [PathNotFoundException] should be thrownBy acyclicGraph.path("e","a")
+  }
 
-  a [CyclicException] should be thrownBy cyclicGraph.linearize
+  "The first element in a linearized graph with a single root node" should "be the root" in {
+    acyclicGraph.linearize.head should equal ("a")
+  }
 
-  try {
-    cyclicGraph.linearize
+  "A DiGraph with a cycle" should "error when linearized" in {
+    a [CyclicException] should be thrownBy cyclicGraph.linearize
   }
-  catch {
-    case c: CyclicException =>
-      c.getMessage.contains("found at a") should be (true)
-      c.node.asInstanceOf[String] should be ("a")
-    case _: Throwable =>
+
+  "CyclicExceptions" should "contain information about the cycle" in {
+    val c = the [CyclicException] thrownBy {
+      cyclicGraph.linearize
+    }
+    c.getMessage.contains("found at a") should be (true)
+    c.node.asInstanceOf[String] should be ("a")
   }
 
-  acyclicGraph.reverse.getEdgeMap should equal (reversedAcyclicGraph.getEdgeMap)
+  "Reversing a graph" should "reverse all of the edges" in {
+    acyclicGraph.reverse.getEdgeMap should equal (reversedAcyclicGraph.getEdgeMap)
+  }
 
-  degenerateGraph.getEdgeMap should equal (degenerateGraph.reverse.getEdgeMap)
+  "Reversing a graph with no edges" should "equal the graph itself" in {
+    degenerateGraph.getEdgeMap should equal (degenerateGraph.reverse.getEdgeMap)
+  }
 
   "transformNodes" should "combine vertices that collide, not drop them" in {
     tupleGraph.transformNodes(_._1).getEdgeMap should contain ("a" -> Set("b", "c"))
@@ -84,4 +98,16 @@ class DiGraphTests extends FirrtlFlatSpec {
     (first + second + second + second).getEdgeMap should equal (acyclicGraph.getEdgeMap)
   }
 
+  "linearize" should "not cause a stack overflow on very large graphs" in {
+    // Graph of 0 -> 1, 1 -> 2, etc.
+    val N = 10000
+    val edges = (1 to N).zipWithIndex.map({ case (n, idx) => idx -> Set(n)}).toMap
+    val bigGraph = DiGraph(edges + (N -> Set.empty[Int]))
+    bigGraph.linearize should be (0 to N)
+  }
+
+  it should "work on multi-rooted graphs" in {
+    val graph = DiGraph(Map("a" -> Set[String](), "b" -> Set[String]()))
+    graph.linearize.toSet should be (graph.getVertices)
+  }
 }