Make `-e` option work with ChiselStage methods (#1630)

* Fix `-e` option causing ChiselStage.emit* to error

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

* Add test of `-e` ChiselStage behavior

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

* fixup! Add test of `-e` ChiselStage behavior

2 files changed, 38 insertions, 13 deletions

diff --git a/src/main/scala/chisel3/stage/ChiselStage.scala b/src/main/scala/chisel3/stage/ChiselStage.scala
index 1ab8b766..aae7ad8d 100644
--- a/src/main/scala/chisel3/stage/ChiselStage.scala
+++ b/src/main/scala/chisel3/stage/ChiselStage.scala
@@ -6,7 +6,9 @@ import firrtl.{
   ir => fir,
   AnnotationSeq,
   EmittedFirrtlCircuitAnnotation,
+  EmittedFirrtlModuleAnnotation,
   EmittedVerilogCircuitAnnotation,
+  EmittedVerilogModuleAnnotation,
   HighFirrtlEmitter,
   VerilogEmitter,
   SystemVerilogEmitter
@@ -93,11 +95,11 @@ class ChiselStage extends Stage {
     annotations: AnnotationSeq = Seq.empty): String = {
 
     execute(Array("-X", "high") ++ args, ChiselGeneratorAnnotation(() => gen) +: annotations)
-      .collectFirst {
+      .collect {
         case EmittedFirrtlCircuitAnnotation(a) => a
-      }
-      .get
-      .value
+        case EmittedFirrtlModuleAnnotation(a)  => a
+      }.map(_.value)
+      .mkString("")
 
   }
 
@@ -115,9 +117,10 @@ class ChiselStage extends Stage {
     execute(Array("-X", "verilog") ++ args, ChiselGeneratorAnnotation(() => gen) +: annotations)
       .collectFirst {
         case EmittedVerilogCircuitAnnotation(a) => a
-      }
-      .get
-      .value
+        case EmittedVerilogModuleAnnotation(a)  => a
+      }.map(_.value)
+      .mkString("")
+
   }
 
   /** Convert a Chisel module to SystemVerilog
@@ -134,9 +137,10 @@ class ChiselStage extends Stage {
     execute(Array("-X", "sverilog") ++ args, ChiselGeneratorAnnotation(() => gen) +: annotations)
       .collectFirst {
         case EmittedVerilogCircuitAnnotation(a) => a
-      }
-      .get
-      .value
+        case EmittedVerilogModuleAnnotation(a)  => a
+      }.map(_.value)
+      .mkString("")
+
   }
 }
 
diff --git a/src/test/scala/chiselTests/stage/ChiselStageSpec.scala b/src/test/scala/chiselTests/stage/ChiselStageSpec.scala
index 09fb3429..e640def8 100644
--- a/src/test/scala/chiselTests/stage/ChiselStageSpec.scala
+++ b/src/test/scala/chiselTests/stage/ChiselStageSpec.scala
@@ -4,6 +4,7 @@ package chiselTests.stage
 
 import chisel3._
 import chisel3.stage.ChiselStage
+import chisel3.testers.TesterDriver.createTestDirectory
 
 import chiselTests.Utils
 
@@ -14,11 +15,19 @@ import firrtl.options.Dependency
 
 object ChiselStageSpec {
 
+  class Bar extends MultiIOModule {
+    val in = IO(Input(UInt(4.W)))
+    val out = IO(Output(UInt(4.W)))
+    out := ~in
+  }
+
   class Foo extends MultiIOModule {
     val addr = IO(Input(UInt(4.W)))
     val out = IO(Output(Bool()))
-    val bar = SyncReadMem(8, Bool())
-    out := bar(addr)
+    val memory = SyncReadMem(8, Bool())
+    val bar = Module(new Bar)
+    bar.in := addr
+    out := memory(bar.out)
   }
 
 }
@@ -40,7 +49,13 @@ class ChiselStageSpec extends AnyFlatSpec with Matchers with Utils {
   behavior of "ChiselStage.emitFirrtl"
 
   it should "return a High FIRRTL string" in {
-    ChiselStage.emitFirrtl(new Foo) should include ("mem bar")
+    ChiselStage.emitFirrtl(new Foo) should include ("mem memory")
+  }
+
+  it should "return a flattened FIRRTL string with '-e high'" in {
+    val args = Array("-e", "high", "-td", createTestDirectory(this.getClass.getSimpleName).toString)
+    (new ChiselStage)
+      .emitFirrtl(new Foo, args) should include ("module Bar")
   }
 
   behavior of "ChiselStage.emitVerilog"
@@ -49,6 +64,12 @@ class ChiselStageSpec extends AnyFlatSpec with Matchers with Utils {
     ChiselStage.emitVerilog(new Foo) should include ("endmodule")
   }
 
+  it should "return a flattened Verilog string with '-e verilog'" in {
+    val args = Array("-e", "verilog", "-td", createTestDirectory(this.getClass.getSimpleName).toString)
+    (new ChiselStage)
+      .emitVerilog(new Foo, args) should include ("module Bar")
+  }
+
   behavior of "ChiselStage$.elaborate"
 
   it should "generate a Chisel circuit from a Chisel module" in {