Merge pull request #1387 from freechipsproject/InstanceGraph.staticInstanceCount-fix

Include Dead Modules in InstanceGraph.staticInstanceCount

2 files changed, 63 insertions, 8 deletions

diff --git a/src/main/scala/firrtl/analyses/InstanceGraph.scala b/src/main/scala/firrtl/analyses/InstanceGraph.scala
index 1a453a42..b8b40065 100644
--- a/src/main/scala/firrtl/analyses/InstanceGraph.scala
+++ b/src/main/scala/firrtl/analyses/InstanceGraph.scala
@@ -62,16 +62,22 @@ class InstanceGraph(c: Circuit) {
     */
   lazy val fullHierarchy: mutable.LinkedHashMap[WDefInstance,Seq[Seq[WDefInstance]]] = graph.pathsInDAG(trueTopInstance)
 
-  /** A count of the *static* number of instances of each module. For any module
-    * other than the top (main) module, this is equivalent to the number of inst
-    * statements in the circuit instantiating each module, irrespective of the
-    * number of times (if any) the enclosing module appears in the hierarchy.
-    * Note that top module of the circuit has an associated count of 1, even
-    * though it is never directly instantiated.
+  /** A count of the *static* number of instances of each module. For any module other than the top (main) module, this is
+    * equivalent to the number of inst statements in the circuit instantiating each module, irrespective of the number
+    * of times (if any) the enclosing module appears in the hierarchy. Note that top module of the circuit has an
+    * associated count of one, even though it is never directly instantiated. Any modules *not* instantiated at all will
+    * have a count of zero.
     */
   lazy val staticInstanceCount: Map[OfModule, Int] = {
-    val instModules = childInstances.flatMap(_._2.view.map(_.OfModule).toSeq)
-    instModules.foldLeft(Map(c.main.OfModule -> 1)) { case (counts, mod) => counts.updated(mod, counts.getOrElse(mod, 0) + 1) }
+    val foo = mutable.LinkedHashMap.empty[OfModule, Int]
+    childInstances.keys.foreach {
+      case main if main == c.main => foo += main.OfModule  -> 1
+      case other                  => foo += other.OfModule -> 0
+    }
+    childInstances.values.flatten.map(_.OfModule).foreach {
+      case mod => foo += mod -> (foo(mod) + 1)
+    }
+    foo.toMap
   }
 
   /** Finds the absolute paths (each represented by a Seq of instances
diff --git a/src/test/scala/firrtlTests/analyses/InstanceGraphTests.scala b/src/test/scala/firrtlTests/analyses/InstanceGraphTests.scala
index eb62c564..ee6ecd5f 100644
--- a/src/test/scala/firrtlTests/analyses/InstanceGraphTests.scala
+++ b/src/test/scala/firrtlTests/analyses/InstanceGraphTests.scala
@@ -1,5 +1,6 @@
 package firrtlTests.analyses
 
+import firrtl.annotations.TargetToken.OfModule
 import firrtl.analyses.InstanceGraph
 import firrtl.graph.DiGraph
 import firrtl.WDefInstance
@@ -195,4 +196,52 @@ circuit Top :
     val hier = instGraph.fullHierarchy
     hier.keys.toSeq.map(_.name) should equal (Seq("Top", "a", "b", "c", "d", "e"))
   }
+
+  behavior of "InstanceGraph.staticInstanceCount"
+
+  it should "report that there is one instance of the top module" in {
+    val input =
+      """|circuit Foo:
+         |  module Foo:
+         |    skip
+         |""".stripMargin
+    val iGraph = new InstanceGraph(ToWorkingIR.run(parse(input)))
+    val expectedCounts = Map(OfModule("Foo") -> 1)
+    iGraph.staticInstanceCount should be (expectedCounts)
+  }
+
+  it should "report correct number of instances for a sample circuit" in {
+    val input =
+      """|circuit Foo:
+         |  module Baz:
+         |    skip
+         |  module Bar:
+         |    inst baz1 of Baz
+         |    inst baz2 of Baz
+         |    inst baz3 of Baz
+         |    skip
+         |  module Foo:
+         |    inst bar1 of Bar
+         |    inst bar2 of Bar
+         |""".stripMargin
+    val iGraph = new InstanceGraph(ToWorkingIR.run(parse(input)))
+    val expectedCounts = Map(OfModule("Foo") -> 1,
+                             OfModule("Bar") -> 2,
+                             OfModule("Baz") -> 3)
+    iGraph.staticInstanceCount should be (expectedCounts)
+  }
+
+  it should "report zero instances for dead modules" in {
+    val input =
+      """|circuit Foo:
+         |  module Bar:
+         |    skip
+         |  module Foo:
+         |    skip
+         |""".stripMargin
+    val iGraph = new InstanceGraph(ToWorkingIR.run(parse(input)))
+    val expectedCounts = Map(OfModule("Foo") -> 1,
+                             OfModule("Bar") -> 0)
+    iGraph.staticInstanceCount should be (expectedCounts)
+  }
 }