bug fix for cases when we want to flatten a module in which a module is instantiated multiple times (#634)

2 files changed, 42 insertions, 3 deletions

diff --git a/src/main/scala/firrtl/transforms/Flatten.scala b/src/main/scala/firrtl/transforms/Flatten.scala
index 748ea00c..cc40c569 100644
--- a/src/main/scala/firrtl/transforms/Flatten.scala
+++ b/src/main/scala/firrtl/transforms/Flatten.scala
@@ -65,7 +65,7 @@ class Flatten extends Transform {
        case WDefInstance(info, instName, moduleName, instTpe) =>
          if (insts.contains(ComponentName(instName, ModuleName(parent.name, CircuitName(c.main))))
            || mods.contains(ModuleName(parent.name, CircuitName(c.main)))) {
-           val newModName = nsp.newName(moduleName+"_TO_FLATTEN")
+           val newModName = if (seedMods.contains(moduleName)) seedMods(moduleName) else nsp.newName(moduleName+"_TO_FLATTEN")
            seedMods += moduleName -> newModName
            WDefInstance(info, instName, newModName, instTpe)
          } else x
@@ -85,7 +85,7 @@ class Flatten extends Transform {
        def dupMod(x: Statement): Statement = x match {
          case _: Block => x map dupMod
          case WDefInstance(info, instName, moduleName, instTpe) =>
-           val newModName = nsp.newName(moduleName+"_TO_FLATTEN")
+           val newModName = if (replMods.contains(moduleName)) replMods(moduleName) else nsp.newName(moduleName+"_TO_FLATTEN")
            replMods += moduleName -> newModName
            WDefInstance(info, instName, newModName, instTpe)
          case _ => x 
diff --git a/src/test/scala/firrtlTests/FlattenTests.scala b/src/test/scala/firrtlTests/FlattenTests.scala
index f695cf4a..10988f8f 100644
--- a/src/test/scala/firrtlTests/FlattenTests.scala
+++ b/src/test/scala/firrtlTests/FlattenTests.scala
@@ -26,7 +26,6 @@ class FlattenTests extends LowTransformSpec {
     FlattenAnnotation(name)
   }
   
-  
   "The modules inside Top " should "be inlined" in {
      val input =
         """circuit Top :
@@ -56,6 +55,46 @@ class FlattenTests extends LowTransformSpec {
           |    b <= a""".stripMargin
      execute(input, check, Seq(flatten("Top")))
   }
+  
+  "Two instances of the same module inside Top " should "be inlined" in {
+     val input =
+        """circuit Top :
+          |  module Top :
+          |    input a : UInt<32>
+          |    output b : UInt<32>
+          |    inst i1 of Inline1
+          |    inst i2 of Inline1
+          |    wire tmp : UInt<32>
+          |    i1.a <= a
+          |    tmp <= i1.b
+          |    i2.a <= tmp
+          |    b <= i2.b
+          |  module Inline1 :
+          |    input a : UInt<32>
+          |    output b : UInt<32>
+          |    b <= a""".stripMargin
+     val check =
+        """circuit Top :
+          |  module Top :
+          |    input a : UInt<32>
+          |    output b : UInt<32>
+          |    wire i1$a : UInt<32>
+          |    wire i1$b : UInt<32>
+          |    i1$b <= i1$a
+          |    wire i2$a : UInt<32>
+          |    wire i2$b : UInt<32>
+          |    i2$b <= i2$a
+          |    wire tmp : UInt<32>
+          |    b <= i2$b
+          |    tmp <= i1$b
+          |    i1$a <= a
+          |    i2$a <= tmp
+          |  module Inline1 :
+          |    input a : UInt<32>
+          |    output b : UInt<32>
+          |    b <= a""".stripMargin
+     execute(input, check, Seq(flatten("Top")))
+  }
 
   "The module instance i in Top " should "be inlined" in {
      val input =