Fix bug incorrectly propagating constants on submodule inputs

Two instances of the same module could collide in counting the number of
instances of each Module. This could lead to constants being propagated
on inputs when it is incorrect to do so.

Fixes #734

2 files changed, 12 insertions, 7 deletions

diff --git a/src/main/scala/firrtl/transforms/ConstantPropagation.scala b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
index ca48cbb5..b767ef65 100644
--- a/src/main/scala/firrtl/transforms/ConstantPropagation.scala
+++ b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
@@ -407,8 +407,13 @@ class ConstantPropagation extends Transform {
       mod.module -> children.map(i => i.name -> i.module).toMap
     })
 
-    // Module name to number of instances
-    val instCount: Map[String, Int] = iGraph.getVertices.groupBy(_.module).mapValues(_.size)
+    // This is a *relative* instance count, ie. how many there are when you visit each Module once
+    // (even if it is instantiated multiple times)
+    val instCount: Map[String, Int] = iGraph.getEdgeMap.foldLeft(Map(c.main -> 1)) {
+      case (cs, (_, values)) => values.foldLeft(cs) {
+        case (counts, value) => counts.updated(value.module, counts.getOrElse(value.module, 0) + 1)
+      }
+    }
 
     // DiGraph using Module names as nodes, destination of edge is a parent Module
     val parentGraph: DiGraph[String] = iGraph.reverse.transformNodes(_.module)
diff --git a/src/test/scala/firrtlTests/ConstantPropagationTests.scala b/src/test/scala/firrtlTests/ConstantPropagationTests.scala
index 06e24b97..2aad4f21 100644
--- a/src/test/scala/firrtlTests/ConstantPropagationTests.scala
+++ b/src/test/scala/firrtlTests/ConstantPropagationTests.scala
@@ -67,9 +67,9 @@ s"""circuit Top :
     out <= in
   module Child :
     output out : UInt<1>
-    inst b of Bottom
-    b.in <= UInt(1)
-    out <= b.out
+    inst b0 of Bottom
+    b0.in <= UInt(1)
+    out <= b0.out
   module Top :
     input x : UInt<1>
     output z : UInt<1>
@@ -91,8 +91,8 @@ s"""circuit Top :
     out <= UInt(1)
   module Child :
     output out : UInt<1>
-    inst b of Bottom
-    b.in <= UInt(1)
+    inst b0 of Bottom
+    b0.in <= UInt(1)
     out <= UInt(1)
   module Top :
     input x : UInt<1>