Resolve register clock dependencies in RemoveWires (#823)

Candidate fix for #749
This adds DefRegister netlist ordering to RemoveWires

Signed-off-by: Schuyler Eldridge <schuyler.eldridge@ibm.com>

2 files changed, 43 insertions, 4 deletions

diff --git a/src/main/scala/firrtl/transforms/RemoveWires.scala b/src/main/scala/firrtl/transforms/RemoveWires.scala
index 5ba953cd..8cc967dd 100644
--- a/src/main/scala/firrtl/transforms/RemoveWires.scala
+++ b/src/main/scala/firrtl/transforms/RemoveWires.scala
@@ -40,7 +40,8 @@ class RemoveWires extends Transform {
 
   // Transform netlist into DefNodes
   private def getOrderedNodes(
-      netlist: mutable.LinkedHashMap[WrappedExpression, (Expression, Info)]): Try[Seq[DefNode]] = {
+    netlist: mutable.LinkedHashMap[WrappedExpression, (Expression, Info)],
+    regInfo: mutable.Map[WrappedExpression, DefRegister]): Try[Seq[Statement]] = {
     val digraph = new MutableDiGraph[WrappedExpression]
     for ((sink, (expr, _)) <- netlist) {
       digraph.addVertex(sink)
@@ -55,9 +56,12 @@ class RemoveWires extends Transform {
     Try {
       val ordered = digraph.linearize.reverse
       ordered.map { key =>
-        val WRef(name, _,_,_) = key.e1
+        val WRef(name, _, kind, _) = key.e1
         val (rhs, info) = netlist(key)
-        DefNode(info, name, rhs)
+        kind match {
+          case RegKind => regInfo(key)
+          case _ => DefNode(info, name, rhs)
+        }
       }
     }
   }
@@ -71,6 +75,8 @@ class RemoveWires extends Transform {
     val netlist = mutable.LinkedHashMap.empty[WrappedExpression, (Expression, Info)]
     // Info at definition of wires for combining into node
     val wireInfo = mutable.HashMap.empty[WrappedExpression, Info]
+    // Additional info about registers
+    val regInfo = mutable.HashMap.empty[WrappedExpression, DefRegister]
 
     def onStmt(stmt: Statement): Statement = {
       stmt match {
@@ -78,6 +84,9 @@ class RemoveWires extends Transform {
           netlist(we(WRef(name))) = (expr, info)
         case wire: DefWire if !wire.tpe.isInstanceOf[AnalogType] => // Remove all non-Analog wires
           wireInfo(WRef(wire)) = wire.info
+        case reg: DefRegister =>
+          regInfo(we(WRef(reg))) = reg
+          netlist(we(WRef(reg))) = (reg.clock, reg.info)
         case decl: IsDeclaration => // Keep all declarations except for nodes and non-Analog wires
           decls += decl
         case con @ Connect(cinfo, lhs, rhs) => kind(lhs) match {
@@ -107,7 +116,7 @@ class RemoveWires extends Transform {
     m match {
       case mod @ Module(info, name, ports, body) =>
         onStmt(body)
-        getOrderedNodes(netlist) match {
+        getOrderedNodes(netlist, regInfo) match {
           case Success(logic) =>
             Module(info, name, ports, Block(decls ++ logic ++ otherStmts))
           // If we hit a CyclicException, just abort removing wires
diff --git a/src/test/scala/firrtlTests/RemoveWiresSpec.scala b/src/test/scala/firrtlTests/RemoveWiresSpec.scala
index cfc03ad9..f162f32c 100644
--- a/src/test/scala/firrtlTests/RemoveWiresSpec.scala
+++ b/src/test/scala/firrtlTests/RemoveWiresSpec.scala
@@ -40,6 +40,24 @@ class RemoveWiresSpec extends FirrtlFlatSpec {
     (nodes, wires)
   }
 
+  def orderedNames(circuit: Circuit): Seq[String] = {
+    require(circuit.modules.size == 1)
+    val names = mutable.ArrayBuffer.empty[String]
+    def onStmt(stmt: Statement): Statement = {
+      stmt map onStmt match {
+        case reg: DefRegister => names += reg.name
+        case wire: DefWire => names += wire.name
+        case node: DefNode => names += node.name
+        case _ =>
+      }
+      stmt
+    }
+    circuit.modules.head match {
+      case Module(_,_,_, body) => onStmt(body)
+    }
+    names
+  }
+
   "Remove Wires" should "turn wires and their single connect into nodes" in {
     val result = compileBody(s"""
       |input a : UInt<8>
@@ -120,4 +138,16 @@ class RemoveWiresSpec extends FirrtlFlatSpec {
           "node y = not(b)")
     )
   }
+
+  it should "work for multiple clocks" in {
+    val result = compileBody(
+      s"""|input clock: Clock
+          |reg a : UInt<1>, clock
+          |node clock2 = asClock(a)
+          |reg b : UInt<1>, clock2
+          |""".stripMargin
+    )
+    val names = orderedNames(result.circuit)
+    names should be (Seq("a", "clock2", "b"))
+  }
 }