Add a pass to deduplicate modules

2 files changed, 92 insertions, 2 deletions

diff --git a/src/main/scala/firrtl/Compiler.scala b/src/main/scala/firrtl/Compiler.scala
index 9781972e..949807db 100644
--- a/src/main/scala/firrtl/Compiler.scala
+++ b/src/main/scala/firrtl/Compiler.scala
@@ -176,8 +176,9 @@ object CompilerUtils {
     } else {
       inputForm match {
         case ChirrtlForm => Seq(new ChirrtlToHighFirrtl) ++ getLoweringTransforms(HighForm, outputForm)
-        case HighForm => Seq(new IRToWorkingIR, new ResolveAndCheck, new HighFirrtlToMiddleFirrtl) ++
-                         getLoweringTransforms(MidForm, outputForm)
+        case HighForm =>
+          Seq(new IRToWorkingIR, new ResolveAndCheck, new transforms.DedupModules,
+              new HighFirrtlToMiddleFirrtl) ++ getLoweringTransforms(MidForm, outputForm)
         case MidForm => Seq(new MiddleFirrtlToLowFirrtl) ++ getLoweringTransforms(LowForm, outputForm)
         case LowForm => error("Internal Error! This shouldn't be possible") // should be caught by if above
       }
diff --git a/src/main/scala/firrtl/transforms/Dedup.scala b/src/main/scala/firrtl/transforms/Dedup.scala
new file mode 100644
index 00000000..5d953e73
--- /dev/null
+++ b/src/main/scala/firrtl/transforms/Dedup.scala
@@ -0,0 +1,89 @@
+// See LICENSE for license details.
+
+package firrtl
+package transforms
+
+import firrtl.ir._
+import firrtl.Mappers._
+import firrtl.Annotations._
+import firrtl.passes.PassException
+
+// Datastructures
+import scala.collection.mutable
+
+// Tags an annotation to be consumed by this pass
+case class DedupAnnotation(target: Named) extends Annotation with Loose with Unstable {
+  def duplicate(n: Named) = this.copy(target=n)
+  def transform = classOf[DedupModules]
+}
+
+// Only use on legal Firrtl. Specifically, the restriction of
+//  instance loops must have been checked, or else this pass can
+//  infinitely recurse
+class DedupModules extends Transform {
+  def inputForm = HighForm
+  def outputForm = HighForm
+  def execute(state: CircuitState): CircuitState = state.copy(circuit = run(state.circuit))
+  def run(c: Circuit): Circuit = {
+    val moduleOrder = mutable.ArrayBuffer.empty[String]
+    val moduleMap = c.modules.map(m => m.name -> m).toMap
+    def hasInstance(b: Statement): Boolean = {
+      var has = false
+      def onStmt(s: Statement): Statement = s map onStmt match {
+        case DefInstance(i, n, m) =>
+          if(!(moduleOrder contains m)) has = true
+          s
+        case WDefInstance(i, n, m, t) =>
+          if(!(moduleOrder contains m)) has = true
+          s
+        case _ => s
+      }
+      onStmt(b)
+      has
+    }
+    def addModule(m: DefModule): DefModule = m match {
+      case Module(info, n, ps, b) =>
+        if(!hasInstance(b)) moduleOrder += m.name
+        m
+      case e: ExtModule =>
+        moduleOrder += m.name
+        m
+      case _ => m
+    }
+
+    while((moduleOrder.size < c.modules.size)) {
+      c.modules.foreach(m => if(!moduleOrder.contains(m.name)) addModule(m))
+    }
+
+    // Module body -> Module name
+    val dedupModules = mutable.HashMap.empty[String, String]
+    // Old module name -> dup module name
+    val dedupMap = mutable.HashMap.empty[String, String]
+    def onModule(m: DefModule): Option[DefModule] = {
+      def fixInstance(s: Statement): Statement = s map fixInstance match {
+        case DefInstance(i, n, m) => DefInstance(i, n, dedupMap.getOrElse(m, m))
+        case WDefInstance(i, n, m, t) => WDefInstance(i, n, dedupMap.getOrElse(m, m), t)
+        case x => x
+      }
+
+      val mx = m map fixInstance
+      val string = mx match {
+        case Module(i, n, ps, b) =>
+          ps.map(_.serialize).mkString + b.serialize
+        case ExtModule(i, n, ps, dn, p) =>
+          ps.map(_.serialize).mkString + dn + p.map(_.serialize).mkString
+      }
+      dedupModules.get(string) match {
+        case Some(dupname) =>
+          dedupMap(mx.name) = dupname
+          None
+        case None =>
+          dedupModules(string) = mx.name
+          Some(mx)
+      }
+    }
+    val modulesx = moduleOrder.flatMap(n => onModule(moduleMap(n)))
+    val modulesxMap = modulesx.map(m => m.name -> m).toMap
+    c.copy(modules = c.modules.flatMap(m => modulesxMap.get(m.name)))
+  }
+}