1 files changed, 43 insertions, 9 deletions

diff --git a/src/main/scala/firrtl/passes/ZeroWidth.scala b/src/main/scala/firrtl/passes/ZeroWidth.scala
index a2ec9935..2472a0e5 100644
--- a/src/main/scala/firrtl/passes/ZeroWidth.scala
+++ b/src/main/scala/firrtl/passes/ZeroWidth.scala
@@ -10,8 +10,24 @@ import firrtl.Mappers._
 import firrtl.Utils.throwInternalError
 
 
-object ZeroWidth extends Pass {
+object ZeroWidth extends Transform {
+  def inputForm = UnknownForm
+  def outputForm = UnknownForm
   private val ZERO = BigInt(0)
+  private def getRemoved(x: IsDeclaration): Seq[String] = {
+    var removedNames: Seq[String] = Seq.empty
+    def onType(name: String)(t: Type): Type = {
+      removedNames = Utils.create_exps(name, t) map {e => (e, e.tpe)} collect {
+        case (e, GroundType(IntWidth(ZERO))) => e.serialize
+      }
+      t
+    }
+    x match {
+      case s: Statement => s map onType(s.name)
+      case Port(_, name, _, t) => onType(name)(t)
+    }
+    removedNames
+  }
   private def removeZero(t: Type): Option[Type] = t match {
     case GroundType(IntWidth(ZERO)) => None
     case BundleType(fields) =>
@@ -34,35 +50,53 @@ object ZeroWidth extends Pass {
       def replaceType(x: Type): Type = t
       (e map replaceType) map onExp
   }
-  private def onStmt(s: Statement): Statement = s match {
+  private def onStmt(renames: RenameMap)(s: Statement): Statement = s match {
     case (_: DefWire| _: DefRegister| _: DefMemory) =>
+      // List all removed expression names, and delete them from renames
+      renames.delete(getRemoved(s.asInstanceOf[IsDeclaration]))
+      // Create new types without zero-width wires
       var removed = false
       def applyRemoveZero(t: Type): Type = removeZero(t) match {
         case None => removed = true; t
         case Some(tx) => tx
       }
       val sxx = (s map onExp) map applyRemoveZero
+      // Return new declaration
       if(removed) EmptyStmt else sxx
     case Connect(info, loc, exp) => removeZero(loc.tpe) match {
       case None => EmptyStmt
       case Some(t) => Connect(info, loc, onExp(exp))
     }
+    case IsInvalid(info, exp) => removeZero(exp.tpe) match {
+      case None => EmptyStmt
+      case Some(t) => IsInvalid(info, onExp(exp))
+    }
     case DefNode(info, name, value) => removeZero(value.tpe) match {
       case None => EmptyStmt
       case Some(t) => DefNode(info, name, onExp(value))
     }
-    case sx => sx map onStmt
+    case sx => sx map onStmt(renames)
   }
-  private def onModule(m: DefModule): DefModule = {
-    val ports = m.ports map (p => (p, removeZero(p.tpe))) collect {
-      case (Port(info, name, dir, _), Some(t)) => Port(info, name, dir, t)
+  private def onModule(renames: RenameMap)(m: DefModule): DefModule = {
+    renames.setModule(m.name)
+    // For each port, record deleted subcomponents
+    m.ports.foreach{p => renames.delete(getRemoved(p))}
+    val ports = m.ports map (p => (p, removeZero(p.tpe))) flatMap {
+      case (Port(info, name, dir, _), Some(t)) => Seq(Port(info, name, dir, t))
+      case (Port(_, name, _, _), None) =>
+        renames.delete(name)
+        Nil
     }
     m match {
       case ext: ExtModule => ext.copy(ports = ports)
-      case in: Module => in.copy(ports = ports, body = onStmt(in.body))
+      case in: Module => in.copy(ports = ports, body = onStmt(renames)(in.body))
     }
   }
-  def run(c: Circuit): Circuit = {
-    InferTypes.run(c.copy(modules = c.modules map onModule))
+  def execute(state: CircuitState): CircuitState = {
+    val c = state.circuit
+    val renames = RenameMap()
+    renames.setCircuit(c.main)
+    val result = InferTypes.run(c.copy(modules = c.modules map onModule(renames)))
+    CircuitState(result, outputForm, state.annotations, Some(renames))
   }
 }