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.

6 files changed, 82 insertions, 8 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
+  }
+}