Fix bug in DiGraph.reverse on an graph with one vertex, no edges (#628)

2 files changed, 16 insertions, 1 deletions

diff --git a/src/main/scala/firrtl/graph/DiGraph.scala b/src/main/scala/firrtl/graph/DiGraph.scala
index e28e53e5..ee00789e 100644
--- a/src/main/scala/firrtl/graph/DiGraph.scala
+++ b/src/main/scala/firrtl/graph/DiGraph.scala
@@ -295,7 +295,10 @@ class DiGraph[T] (val edges: Map[T, Set[T]]) extends DiGraphLike[T] {
   /** Returns a graph with all edges reversed */
   def reverse: DiGraph[T] = {
     val mdg = new MutableDiGraph[T]
-    edges foreach { case (u,edges) => edges.foreach({ v => mdg.addEdge(v,u) }) }
+    edges.foreach { case (u, edges) =>
+      mdg.addVertex(u)
+      edges.foreach(v => mdg.addEdge(v,u))
+    }
     DiGraph(mdg)
   }
 
diff --git a/src/test/scala/firrtlTests/graph/DiGraphTests.scala b/src/test/scala/firrtlTests/graph/DiGraphTests.scala
index 6546e147..9eb1c7f8 100644
--- a/src/test/scala/firrtlTests/graph/DiGraphTests.scala
+++ b/src/test/scala/firrtlTests/graph/DiGraphTests.scala
@@ -16,12 +16,20 @@ class DiGraphTests extends FirrtlFlatSpec {
     "d" -> Set("e"),
     "e" -> Set.empty[String]))
 
+  val reversedAcyclicGraph = DiGraph(Map(
+    "a" -> Set.empty[String],
+    "b" -> Set("a"),
+    "c" -> Set("a"),
+    "d" -> Set("b", "c"),
+    "e" -> Set("d")))
+
   val cyclicGraph = DiGraph(Map(
     "a" -> Set("b","c"),
     "b" -> Set("d"),
     "c" -> Set("d"),
     "d" -> Set("a")))
 
+  val degenerateGraph = DiGraph(Map("a" -> Set.empty[String]))
 
   acyclicGraph.findSCCs.filter(_.length > 1) shouldBe empty
 
@@ -35,4 +43,8 @@ class DiGraphTests extends FirrtlFlatSpec {
 
   a [cyclicGraph.CyclicException] should be thrownBy cyclicGraph.linearize
 
+  acyclicGraph.reverse.edges should equal (reversedAcyclicGraph.edges)
+
+  degenerateGraph.edges should equal (degenerateGraph.reverse.edges)
+
 }