Make submodule inputs void in ExpandWhens (#706)

4 files changed, 39 insertions, 1 deletions

diff --git a/src/main/scala/firrtl/passes/ExpandWhens.scala b/src/main/scala/firrtl/passes/ExpandWhens.scala
index 181fb642..6ff0debe 100644
--- a/src/main/scala/firrtl/passes/ExpandWhens.scala
+++ b/src/main/scala/firrtl/passes/ExpandWhens.scala
@@ -103,12 +103,16 @@ object ExpandWhens extends Pass {
                       defaults: Defaults,
                       p: Expression)
                       (s: Statement): Statement = s match {
+        case stmt @ (_: DefNode | EmptyStmt) => stmt
         case w: DefWire =>
           netlist ++= (getFemaleRefs(w.name, w.tpe, BIGENDER) map (ref => we(ref) -> WVoid))
           w
         case w: DefMemory =>
           netlist ++= (getFemaleRefs(w.name, MemPortUtils.memType(w), MALE) map (ref => we(ref) -> WVoid))
           w
+        case w: WDefInstance =>
+          netlist ++= (getFemaleRefs(w.name, w.tpe, MALE).map(ref => we(ref) -> WVoid))
+          w
         case r: DefRegister =>
           netlist ++= (getFemaleRefs(r.name, r.tpe, BIGENDER) map (ref => we(ref) -> ref))
           r
@@ -173,7 +177,8 @@ object ExpandWhens extends Pass {
         case sx: Stop =>
           simlist += (if (weq(p, one)) sx else Stop(sx.info, sx.ret, sx.clk, AND(p, sx.en)))
           EmptyStmt
-        case sx => sx map expandWhens(netlist, defaults, p)
+        case block: Block => block map expandWhens(netlist, defaults, p)
+        case _ => throwInternalError
       }
       val netlist = new Netlist
       // Add ports to netlist
diff --git a/src/test/scala/firrtlTests/CheckInitializationSpec.scala b/src/test/scala/firrtlTests/CheckInitializationSpec.scala
index cd6f464d..ef966ca0 100644
--- a/src/test/scala/firrtlTests/CheckInitializationSpec.scala
+++ b/src/test/scala/firrtlTests/CheckInitializationSpec.scala
@@ -55,4 +55,20 @@ class CheckInitializationSpec extends FirrtlFlatSpec {
       }
     }
   }
+  "Missing assignment to submodule port" should "trigger a PassException" in {
+    val input =
+      """circuit Test :
+        |  module Child :
+        |    input in : UInt<32>
+        |  module Test :
+        |    input p : UInt<1>
+        |    inst c of Child
+        |    when p :
+        |      c.in <= UInt(1)""".stripMargin
+    intercept[CheckInitialization.RefNotInitializedException] {
+      passes.foldLeft(parse(input)) {
+        (c: Circuit, p: Pass) => p.run(c)
+      }
+    }
+  }
 }
diff --git a/src/test/scala/firrtlTests/DCETests.scala b/src/test/scala/firrtlTests/DCETests.scala
index d1848ab8..e28ab432 100644
--- a/src/test/scala/firrtlTests/DCETests.scala
+++ b/src/test/scala/firrtlTests/DCETests.scala
@@ -117,6 +117,7 @@ class DCETests extends FirrtlFlatSpec {
         |    output z : UInt<1>
         |    inst sub of Sub
         |    sub.x <= x
+        |    sub.y is invalid
         |    z <= sub.z""".stripMargin
     val check =
       """circuit Top :
diff --git a/src/test/scala/firrtlTests/ExpandWhensSpec.scala b/src/test/scala/firrtlTests/ExpandWhensSpec.scala
index 66f39a3d..3532ce00 100644
--- a/src/test/scala/firrtlTests/ExpandWhensSpec.scala
+++ b/src/test/scala/firrtlTests/ExpandWhensSpec.scala
@@ -122,6 +122,22 @@ class ExpandWhensSpec extends FirrtlFlatSpec {
     val check = "w is invalid"
     executeTest(input, check, false)
   }
+  it should "correctly handle submodule inputs" in {
+    val input =
+      """circuit Test :
+        |  module Child :
+        |    input in : UInt<32>
+        |  module Test :
+        |    input in : UInt<32>[2]
+        |    input p : UInt<1>
+        |    inst c of Child
+        |    when p :
+        |      c.in <= in[0]
+        |    else :
+        |      c.in <= in[1]""".stripMargin
+    val check = "mux(p, in[0], in[1])"
+    executeTest(input, check, true)
+  }
 }
 
 class ExpandWhensExecutionTest extends ExecutionTest("ExpandWhens", "/passes/ExpandWhens")