Add new annotation for Chisel Circuit serialization (#1580)

ChiselCircuitAnnotation no longer extends CustomFileEmission, rather it
is Unserializable. Also the --chisel-output-file is added to the
ChiselCli.

New phase AddSerializationAnnotations constructs a
CircuitSerializationAnnotation from ChiselCircuitAnnotation and
ChiselOutputFileAnnotation. Both .fir and .pb file formats are
supported. Default format is .fir unless a --chisel-output-file is
specified with a .pb extension.

8 files changed, 177 insertions, 10 deletions

diff --git a/src/main/scala/chisel3/stage/ChiselAnnotations.scala b/src/main/scala/chisel3/stage/ChiselAnnotations.scala
index e4663dd4..0b447413 100644
--- a/src/main/scala/chisel3/stage/ChiselAnnotations.scala
+++ b/src/main/scala/chisel3/stage/ChiselAnnotations.scala
@@ -103,20 +103,53 @@ object ChiselGeneratorAnnotation extends HasShellOptions {
 case class ChiselCircuitAnnotation(circuit: Circuit)
     extends NoTargetAnnotation
     with ChiselOption
-    with CustomFileEmission {
+    with Unserializable {
   /* Caching the hashCode for a large circuit is necessary due to repeated queries.
    * Not caching the hashCode will cause severe performance degredations for large [[Circuit]]s.
    */
   override lazy val hashCode: Int = circuit.hashCode
+}
 
-  protected def baseFileName(annotations: AnnotationSeq): String = {
-    view[ChiselOptions](annotations).outputFile.getOrElse(circuit.name)
+object CircuitSerializationAnnotation {
+  sealed trait Format {
+    def extension: String
+  }
+  case object FirrtlFileFormat extends Format {
+    def extension = ".fir"
+  }
+  case object ProtoBufFileFormat extends Format {
+    def extension = ".pb"
   }
+}
 
-  protected def suffix: Option[String] = Some(".fir")
+import CircuitSerializationAnnotation._
 
-  override def getBytes: Iterable[Byte] = OldEmitter.emit(circuit).getBytes
+/** Wraps a [[Circuit]] for serialization via [[CustomFileEmission]]
+  * @param circuit a Chisel Circuit
+  * @param filename name of destination file (excludes file extension)
+  * @param format serialization file format (sets file extension)
+  */
+case class CircuitSerializationAnnotation(circuit: Circuit, filename: String, format: Format)
+    extends NoTargetAnnotation
+    with CustomFileEmission {
+  /* Caching the hashCode for a large circuit is necessary due to repeated queries.
+   * Not caching the hashCode will cause severe performance degredations for large [[Circuit]]s.
+   */
+  override lazy val hashCode: Int = circuit.hashCode
+
+  protected def baseFileName(annotations: AnnotationSeq): String = filename
 
+  protected def suffix: Option[String] = Some(format.extension)
+
+  // TODO Use lazy Iterables so that we don't have to materialize full intermediate data structures
+  override def getBytes: Iterable[Byte] = format match {
+    case FirrtlFileFormat => OldEmitter.emit(circuit).getBytes
+    case ProtoBufFileFormat =>
+      val ostream = new java.io.ByteArrayOutputStream
+      val modules = circuit.components.map(m => () => chisel3.internal.firrtl.Converter.convert(m))
+      firrtl.proto.ToProto.writeToStreamFast(ostream, firrtl.ir.NoInfo, modules, circuit.name)
+      ostream.toByteArray
+  }
 }
 
 case class ChiselOutputFileAnnotation(file: String) extends NoTargetAnnotation with ChiselOption with Unserializable
diff --git a/src/main/scala/chisel3/stage/ChiselCli.scala b/src/main/scala/chisel3/stage/ChiselCli.scala
index 78150029..000c6c71 100644
--- a/src/main/scala/chisel3/stage/ChiselCli.scala
+++ b/src/main/scala/chisel3/stage/ChiselCli.scala
@@ -8,6 +8,7 @@ trait ChiselCli { this: Shell =>
   parser.note("Chisel Front End Options")
   Seq( NoRunFirrtlCompilerAnnotation,
        PrintFullStackTraceAnnotation,
+       ChiselOutputFileAnnotation,
        ChiselGeneratorAnnotation )
     .foreach(_.addOptions(parser))
 }
diff --git a/src/main/scala/chisel3/stage/ChiselPhase.scala b/src/main/scala/chisel3/stage/ChiselPhase.scala
index 200bc2b6..1fbac622 100644
--- a/src/main/scala/chisel3/stage/ChiselPhase.scala
+++ b/src/main/scala/chisel3/stage/ChiselPhase.scala
@@ -22,6 +22,7 @@ private[chisel3] object ChiselPhase {
          Dependency[chisel3.stage.phases.AddImplicitOutputFile],
          Dependency[chisel3.stage.phases.AddImplicitOutputAnnotationFile],
          Dependency[chisel3.stage.phases.MaybeAspectPhase],
+         Dependency[chisel3.stage.phases.AddSerializationAnnotations],
          Dependency[chisel3.stage.phases.Convert],
          Dependency[chisel3.stage.phases.MaybeFirrtlStage] )
 
diff --git a/src/main/scala/chisel3/stage/ChiselStage.scala b/src/main/scala/chisel3/stage/ChiselStage.scala
index 30723430..b9394a88 100644
--- a/src/main/scala/chisel3/stage/ChiselStage.scala
+++ b/src/main/scala/chisel3/stage/ChiselStage.scala
@@ -18,6 +18,7 @@ import firrtl.options.Viewer.view
 
 import chisel3.{ChiselException, RawModule}
 import chisel3.internal.{firrtl => cir, ErrorLog}
+import chisel3.stage.CircuitSerializationAnnotation.FirrtlFileFormat
 
 import java.io.{StringWriter, PrintWriter}
 
@@ -72,7 +73,7 @@ class ChiselStage extends Stage {
 
     annos
       .collectFirst {
-        case a: ChiselCircuitAnnotation => a.getBytes
+        case a: ChiselCircuitAnnotation => CircuitSerializationAnnotation(a.circuit, "", FirrtlFileFormat).getBytes
       }
       .get
       .map(_.toChar)
diff --git a/src/main/scala/chisel3/stage/package.scala b/src/main/scala/chisel3/stage/package.scala
index 079e3f22..5dcc45d6 100644
--- a/src/main/scala/chisel3/stage/package.scala
+++ b/src/main/scala/chisel3/stage/package.scala
@@ -6,6 +6,7 @@ import firrtl._
 import firrtl.options.OptionsView
 
 import chisel3.internal.firrtl.{Circuit => ChiselCircuit}
+import chisel3.stage.CircuitSerializationAnnotation.FirrtlFileFormat
 
 package object stage {
 
@@ -31,9 +32,12 @@ package object stage {
       var chirrtlCircuit: Option[String] = None
 
       options.foreach {
-        case a@ ChiselCircuitAnnotation(b) =>
+        case a @ ChiselCircuitAnnotation(b) =>
           chiselCircuit = Some(b)
-          chirrtlCircuit = Some(a.getBytes.map(_.toChar).mkString)
+          chirrtlCircuit = {
+            val anno = CircuitSerializationAnnotation(a.circuit, "", FirrtlFileFormat)
+            Some(anno.getBytes.map(_.toChar).mkString)
+          }
         case _ =>
       }
 
diff --git a/src/main/scala/chisel3/stage/phases/AddSerializationAnnotations.scala b/src/main/scala/chisel3/stage/phases/AddSerializationAnnotations.scala
new file mode 100644
index 00000000..c8835e07
--- /dev/null
+++ b/src/main/scala/chisel3/stage/phases/AddSerializationAnnotations.scala
@@ -0,0 +1,34 @@
+// See LICENSE for license details.
+
+package chisel3.stage.phases
+
+import firrtl.AnnotationSeq
+import firrtl.options.{Dependency, Phase}
+import firrtl.options.Viewer.view
+
+import chisel3.stage._
+import chisel3.stage.CircuitSerializationAnnotation._
+import chisel3.internal.ChiselException
+
+/** Adds [[stage.CircuitSerializationAnnotation]]s based on [[ChiselOutputFileAnnotation]]
+  */
+class AddSerializationAnnotations extends Phase {
+
+  override def prerequisites = Seq(Dependency[Elaborate], Dependency[AddImplicitOutputFile])
+  override def optionalPrerequisites = Seq.empty
+  override def optionalPrerequisiteOf = Seq.empty
+  override def invalidates(a: Phase) = false
+
+  def transform(annotations: AnnotationSeq): AnnotationSeq = {
+    val chiselOptions = view[ChiselOptions](annotations)
+    val circuit = chiselOptions.chiselCircuit.getOrElse {
+      throw new ChiselException(s"Unable to locate the elaborated circuit, did ${classOf[Elaborate].getName} run correctly")
+    }
+    val baseFilename = chiselOptions.outputFile.getOrElse(circuit.name)
+
+    val (filename, format) =
+      if (baseFilename.endsWith(".pb")) (baseFilename.stripSuffix(".pb"), ProtoBufFileFormat)
+      else (baseFilename.stripSuffix(".fir"), FirrtlFileFormat)
+    CircuitSerializationAnnotation(circuit, filename, format) +: annotations
+  }
+}
diff --git a/src/test/scala/chiselTests/stage/ChiselMainSpec.scala b/src/test/scala/chiselTests/stage/ChiselMainSpec.scala
index 036e15b1..2e9d928f 100644
--- a/src/test/scala/chiselTests/stage/ChiselMainSpec.scala
+++ b/src/test/scala/chiselTests/stage/ChiselMainSpec.scala
@@ -11,6 +11,9 @@ import org.scalatest.GivenWhenThen
 import org.scalatest.featurespec.AnyFeatureSpec
 import org.scalatest.matchers.should.Matchers
 
+import scala.io.Source
+import firrtl.Parser
+
 object ChiselMainSpec {
 
   /** A module that connects two different types together resulting in an elaboration error */
@@ -52,7 +55,7 @@ class ChiselMainSpec extends AnyFeatureSpec with GivenWhenThen with Matchers wit
   }
 
   class TargetDirectoryFixture(dirName: String) {
-    val dir = new File(s"test_run_dir/FirrtlStageSpec/$dirName")
+    val dir = new File(s"test_run_dir/ChiselStageSpec/$dirName")
     val buildDir = new File(dir + "/build")
     dir.mkdirs()
   }
@@ -63,7 +66,8 @@ class ChiselMainSpec extends AnyFeatureSpec with GivenWhenThen with Matchers wit
     files: Seq[String] = Seq.empty,
     stdout: Option[String] = None,
     stderr: Option[String] = None,
-    result: Int = 0) {
+    result: Int = 0,
+    fileChecks: Map[String, File => Unit] = Map.empty) {
     def testName: String = "args" + args.mkString("_")
     def argsString: String = args.mkString(" ")
   }
@@ -117,6 +121,7 @@ class ChiselMainSpec extends AnyFeatureSpec with GivenWhenThen with Matchers wit
         And(s"file '$f' should be emitted in the target directory")
         val out = new File(td.buildDir + s"/$f")
         out should (exist)
+        p.fileChecks.get(f).map(_(out))
       }
     }
   }
@@ -148,6 +153,33 @@ class ChiselMainSpec extends AnyFeatureSpec with GivenWhenThen with Matchers wit
     ).foreach(runStageExpectFiles)
   }
 
+  Feature("Specifying a custom output file") {
+    runStageExpectFiles(ChiselMainTest(
+      args = Array("--chisel-output-file", "Foo", "--no-run-firrtl"),
+      generator = Some(classOf[SameTypesModule]),
+      stdout = Some(""),
+      files = Seq("Foo.fir"),
+      fileChecks = Map(
+        "Foo.fir" -> { file =>
+          And("It should be valid FIRRTL")
+          Parser.parse(Source.fromFile(file).mkString)
+        }
+      )
+    ))
+    runStageExpectFiles(ChiselMainTest(
+      args = Array("--chisel-output-file", "Foo.pb", "--no-run-firrtl"),
+      generator = Some(classOf[SameTypesModule]),
+      stdout = Some(""),
+      files = Seq("Foo.pb"),
+      fileChecks = Map(
+        "Foo.pb" -> { file =>
+          And("It should be valid ProtoBuf")
+          firrtl.proto.FromProto.fromFile(file.toString)
+        }
+      )
+    ))
+  }
+
   info("As an aspect writer")
   info("I write an aspect")
   Feature("Running aspects via the command line") {
diff --git a/src/test/scala/chiselTests/stage/phases/AddSerializationAnnotationsSpec.scala b/src/test/scala/chiselTests/stage/phases/AddSerializationAnnotationsSpec.scala
new file mode 100644
index 00000000..c248da6a
--- /dev/null
+++ b/src/test/scala/chiselTests/stage/phases/AddSerializationAnnotationsSpec.scala
@@ -0,0 +1,61 @@
+// See LICENSE for license details.
+
+package chiselTests.stage.phases
+
+
+import chisel3.RawModule
+import chisel3.stage.{ChiselGeneratorAnnotation, ChiselOutputFileAnnotation, CircuitSerializationAnnotation}
+import chisel3.stage.CircuitSerializationAnnotation._
+import chisel3.stage.phases.{AddSerializationAnnotations, AddImplicitOutputFile, Elaborate}
+
+import firrtl.AnnotationSeq
+import firrtl.options.{Phase, PhaseManager, Dependency, TargetDirAnnotation}
+import firrtl.options.Viewer.view
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+class AddSerializationAnnotationsSpec extends AnyFlatSpec with Matchers {
+
+  class Foo extends RawModule { override val desiredName = "Foo" }
+
+  class Fixture {
+    val phase: Phase = new AddSerializationAnnotations
+    val manager = new PhaseManager(Dependency[AddSerializationAnnotations] :: Nil)
+  }
+
+  behavior of classOf[AddSerializationAnnotations].toString
+
+  it should "default to FirrtlFileFormat" in new Fixture {
+    val annotations: AnnotationSeq = Seq(
+      ChiselGeneratorAnnotation(() => new Foo),
+      ChiselOutputFileAnnotation("Bar") )
+
+    manager
+      .transform(annotations)
+      .collect { case CircuitSerializationAnnotation(_, filename, format) => (filename, format) }
+      .toSeq should be (Seq(("Bar", FirrtlFileFormat)))
+  }
+
+  it should "support ProtoBufFileFormat" in new Fixture {
+    val annotations: AnnotationSeq = Seq(
+      ChiselGeneratorAnnotation(() => new Foo),
+      ChiselOutputFileAnnotation("Bar.pb") )
+
+    manager
+      .transform(annotations)
+      .collect { case CircuitSerializationAnnotation(_, filename, format) => (filename, format) }
+      .toSeq should be (Seq(("Bar", ProtoBufFileFormat)))
+  }
+
+  it should "support explicitly asking for FirrtlFileFormat" in new Fixture {
+    val annotations: AnnotationSeq = Seq(
+      ChiselGeneratorAnnotation(() => new Foo),
+      ChiselOutputFileAnnotation("Bar.pb.fir") )
+
+    manager
+      .transform(annotations)
+      .collect { case CircuitSerializationAnnotation(_, filename, format) => (filename, format) }
+      .toSeq should be (Seq(("Bar.pb", FirrtlFileFormat)))
+  }
+
+}