ExpandWhens: ensure that statement names are maintained (#2082)

2 files changed, 19 insertions, 5 deletions

diff --git a/src/main/scala/firrtl/passes/ExpandWhens.scala b/src/main/scala/firrtl/passes/ExpandWhens.scala
index 7456d2ab..8fb4e5fb 100644
--- a/src/main/scala/firrtl/passes/ExpandWhens.scala
+++ b/src/main/scala/firrtl/passes/ExpandWhens.scala
@@ -125,13 +125,13 @@ object ExpandWhens extends Pass {
         EmptyStmt
       // For simulation constructs, update simlist with predicated statement and return EmptyStmt
       case sx: Print =>
-        simlist += (if (weq(p, one)) sx else Print(sx.info, sx.string, sx.args, sx.clk, AND(p, sx.en)))
+        simlist += (if (weq(p, one)) sx else sx.withEn(AND(p, sx.en)))
         EmptyStmt
       case sx: Stop =>
-        simlist += (if (weq(p, one)) sx else Stop(sx.info, sx.ret, sx.clk, AND(p, sx.en)))
+        simlist += (if (weq(p, one)) sx else sx.withEn(AND(p, sx.en)))
         EmptyStmt
       case sx: Verification =>
-        simlist += (if (weq(p, one)) sx else sx.copy(en = AND(p, sx.en)))
+        simlist += (if (weq(p, one)) sx else sx.withEn(AND(p, sx.en)))
         EmptyStmt
       // Expand conditionally, see comments below
       case sx: Conditionally =>
diff --git a/src/test/scala/firrtlTests/ExpandWhensSpec.scala b/src/test/scala/firrtlTests/ExpandWhensSpec.scala
index c186b516..3c685734 100644
--- a/src/test/scala/firrtlTests/ExpandWhensSpec.scala
+++ b/src/test/scala/firrtlTests/ExpandWhensSpec.scala
@@ -142,10 +142,24 @@ class ExpandWhensSpec extends FirrtlFlatSpec {
         |    input in : UInt<32>
         |    input p : UInt<1>
         |    when p :
-        |      assert(clock, eq(in, UInt<1>("h1")), UInt<1>("h1"), "assert0")
+        |      assert(clock, eq(in, UInt<1>("h1")), UInt<1>("h1"), "assert0") : test_assert
         |    else :
         |      skip""".stripMargin
-    val check = "assert(clock, eq(in, UInt<1>(\"h1\")), and(and(UInt<1>(\"h1\"), p), UInt<1>(\"h1\")), \"assert0\")"
+    val check = "assert(clock, eq(in, UInt<1>(\"h1\")), and(and(UInt<1>(\"h1\"), p), UInt<1>(\"h1\")), \"assert0\") : test_assert"
+    executeTest(input, check, true)
+  }
+  it should "handle stops" in {
+    val input =
+      """circuit Test :
+        |  module Test :
+        |    input clock : Clock
+        |    input in : UInt<32>
+        |    input p : UInt<1>
+        |    when p :
+        |      stop(clock, UInt(1), 1) : test_stop
+        |    else :
+        |      skip""".stripMargin
+    val check = """stop(clock, and(and(UInt<1>("h1"), p), UInt<1>("h1")), 1) : test_stop"""
     executeTest(input, check, true)
   }
 }