Multi protobuf module emission and consumption

src/main/scala/firrtl/Emitter.scala

@@ -99,7 +99,7 @@ object EmitCircuitAnnotation extends HasShellOptions {
           case "low-opt" =>
             Seq(
               RunFirrtlTransformAnnotation(new ProtoEmitter.OptLow),
-              EmitCircuitAnnotation(classOf[ProtoEmitter.Low])
+              EmitCircuitAnnotation(classOf[ProtoEmitter.OptLow])
             )
           case _ => throw new PhaseException(s"Unknown emitter '$a'! (Did you misspell it?)")
         },
@@ -151,6 +151,43 @@ object EmitAllModulesAnnotation extends HasShellOptions {
       helpText = "Run the specified module emitter (one file per module)",
       shortOption = Some("e"),
       helpValueName = Some("<chirrtl|high|middle|low|verilog|mverilog|sverilog>")
+    ),
+    new ShellOption[String](
+      longOption = "emit-module-protobufs",
+      toAnnotationSeq = (a: String) =>
+        a match {
+          case "chirrtl" =>
+            Seq(
+              RunFirrtlTransformAnnotation(new ProtoEmitter.Chirrtl),
+              EmitAllModulesAnnotation(classOf[ProtoEmitter.Chirrtl])
+            )
+          case "mhigh" =>
+            Seq(
+              RunFirrtlTransformAnnotation(new ProtoEmitter.MHigh),
+              EmitAllModulesAnnotation(classOf[ProtoEmitter.MHigh])
+            )
+          case "high" =>
+            Seq(
+              RunFirrtlTransformAnnotation(new ProtoEmitter.High),
+              EmitAllModulesAnnotation(classOf[ProtoEmitter.High])
+            )
+          case "middle" =>
+            Seq(
+              RunFirrtlTransformAnnotation(new ProtoEmitter.Middle),
+              EmitAllModulesAnnotation(classOf[ProtoEmitter.Middle])
+            )
+          case "low" =>
+            Seq(RunFirrtlTransformAnnotation(new ProtoEmitter.Low), EmitAllModulesAnnotation(classOf[ProtoEmitter.Low]))
+          case "low-opt" =>
+            Seq(
+              RunFirrtlTransformAnnotation(new ProtoEmitter.OptLow),
+              EmitAllModulesAnnotation(classOf[ProtoEmitter.OptLow])
+            )
+          case _ => throw new PhaseException(s"Unknown emitter '$a'! (Did you misspell it?)")
+        },
+      helpText = "Run the specified module emitter (one protobuf per module)",
+      shortOption = Some("p"),
+      helpValueName = Some("<chirrtl|mhigh|high|middle|low|low-opt>")
     )
   )
 

src/main/scala/firrtl/Utils.scala

@@ -978,6 +978,37 @@ object Utils extends LazyLogging {
     map.view.map({ case (k, vs) => k -> vs.toList }).toList
   }
 
+  /** Combines several separate circuit modules (typically emitted by -e or -p compiler options) into a single circuit */
+  def combine(circuits: Seq[Circuit]): Circuit = {
+    def dedup(modules: Seq[DefModule]): Seq[Either[Module, DefModule]] = {
+      // Left means "lone module", Right means has ExtModules
+      val module: Option[Module] = {
+        val found: Seq[Module] = modules.collect { case m: Module => m }
+        assert(found.size <= 1)
+        found.headOption
+      }
+      val extModules: Seq[ExtModule] = modules.collect { case e: ExtModule => e }.distinct
+
+      if (extModules.isEmpty) Seq(Left(module.get)) else (module ++: extModules).map(Right(_))
+    }
+
+    // 1. Combine modules
+    val grouped: Seq[(String, Seq[DefModule])] = groupByIntoSeq(circuits.flatMap(_.modules))({
+      case mod: Module    => mod.name
+      case ext: ExtModule => ext.defname
+    })
+    val deduped: Iterable[Either[Module, DefModule]] = grouped.flatMap { case (_, insts) => dedup(insts) }
+
+    // 2. Determine top
+    val top = {
+      val found = deduped.collect { case Left(m) => m }
+      assert(found.size == 1)
+      found.head
+    }
+    val res = deduped.collect { case Right(m: Module) => m }
+    ir.Circuit(NoInfo, top +: res.toSeq, top.name)
+  }
+
   object True {
     private val _True = UIntLiteral(1, IntWidth(1))
 

src/main/scala/firrtl/backends/proto/ProtoBufEmitter.scala

@@ -1,15 +1,18 @@
 // SPDX-License-Identifier: Apache-2.0
 package firrtl.backends.proto
 
-import firrtl.{AnnotationSeq, CircuitState, DependencyAPIMigration, Transform}
-import firrtl.ir
+import firrtl._
+import firrtl.ir._
 import firrtl.annotations.NoTargetAnnotation
 import firrtl.options.CustomFileEmission
 import firrtl.options.Viewer.view
 import firrtl.proto.ToProto
 import firrtl.stage.{FirrtlOptions, Forms}
 import firrtl.stage.TransformManager.TransformDependency
+import firrtl.traversals.Foreachers._
 import java.io.{ByteArrayOutputStream, Writer}
+import scala.collection.mutable.ArrayBuffer
+import Utils.throwInternalError
 
 /** This object defines Annotations that are used by Protocol Buffer emission.
   */
@@ -59,10 +62,47 @@ sealed abstract class ProtoBufEmitter(prereqs: Seq[TransformDependency])
   override def optionalPrerequisiteOf = Seq.empty
   override def invalidates(a: Transform) = false
 
-  override def execute(state: CircuitState) =
-    state.copy(annotations = state.annotations :+ Annotation.ProtoBufSerialization(state.circuit, Some(outputSuffix)))
+  private def emitAllModules(circuit: Circuit): Seq[Annotation.ProtoBufSerialization] = {
+    // For a given module, returns a Seq of all modules instantited inside of it
+    def collectInstantiatedModules(mod: Module, map: Map[String, DefModule]): Seq[DefModule] = {
+      // Use list instead of set to maintain order
+      val modules = ArrayBuffer.empty[DefModule]
+      def onStmt(stmt: Statement): Unit = stmt match {
+        case DefInstance(_, _, name, _) => modules += map(name)
+        case _: WDefInstanceConnector => throwInternalError(s"unrecognized statement: $stmt")
+        case other => other.foreach(onStmt)
+      }
+      onStmt(mod.body)
+      modules.distinct.toSeq
+    }
+    val modMap = circuit.modules.map(m => m.name -> m).toMap
+    // Turn each module into it's own circuit with it as the top and all instantied modules as ExtModules
+    circuit.modules.collect {
+      case m: Module =>
+        val instModules = collectInstantiatedModules(m, modMap)
+        val extModules = instModules.map {
+          case Module(info, name, ports, _) => ExtModule(info, name, ports, name, Seq.empty)
+          case ext: ExtModule => ext
+        }
+        val newCircuit = Circuit(m.info, extModules :+ m, m.name)
+        Annotation.ProtoBufSerialization(newCircuit, Some(outputSuffix))
+    }
+  }
+
+  override def execute(state: CircuitState) = {
+    val newAnnos = state.annotations.flatMap {
+      case EmitCircuitAnnotation(a) if this.getClass == a =>
+        Seq(
+          Annotation.ProtoBufSerialization(state.circuit, Some(outputSuffix))
+        )
+      case EmitAllModulesAnnotation(a) if this.getClass == a =>
+        emitAllModules(state.circuit)
+      case _ => Seq()
+    }
+    state.copy(annotations = newAnnos ++ state.annotations)
+  }
 
-  override def emit(state: CircuitState, writer: Writer): Unit = {
+  def emit(state: CircuitState, writer: Writer): Unit = {
     val ostream = new java.io.ByteArrayOutputStream
     ToProto.writeToStream(ostream, state.circuit)
     writer.write(ostream.toString())

src/main/scala/firrtl/proto/FromProto.scala

@@ -9,6 +9,8 @@ import collection.JavaConverters._
 import FirrtlProtos._
 import com.google.protobuf.CodedInputStream
 import Firrtl.Statement.{Formal, ReadUnderWrite}
+import firrtl.ir.DefModule
+import Utils.combine
 
 object FromProto {
 
@@ -33,6 +35,22 @@ object FromProto {
     proto.FromProto.convert(pb)
   }
 
+  /** Deserialize all the ProtoBuf representations of [[ir.Circuit]] in @directory
+    *
+    * @param dir directory containing ProtoBuf representation(s)
+    * @return Deserialized FIRRTL Circuit
+    */
+  def fromDirectory(dir: String): ir.Circuit = {
+    val d = new File(dir)
+    val fileList = if (d.exists && d.isDirectory) {
+      d.listFiles.filter(_.isFile).toList
+    } else {
+      List[File]()
+    }
+
+    combine(fileList.map(f => fromInputStream(new FileInputStream(f))))
+  }
+
   // Convert from ProtoBuf message repeated Statements to FIRRRTL Block
   private def compressStmts(stmts: scala.collection.Seq[ir.Statement]): ir.Statement = stmts match {
     case scala.collection.Seq()     => ir.EmptyStmt

src/main/scala/firrtl/stage/FirrtlAnnotations.scala

@@ -6,7 +6,7 @@ import firrtl._
 import firrtl.ir.Circuit
 import firrtl.annotations.{Annotation, NoTargetAnnotation}
 import firrtl.options.{Dependency, HasShellOptions, OptionsException, ShellOption, Unserializable}
-import java.io.FileNotFoundException
+import java.io.{File, FileNotFoundException}
 import java.nio.file.NoSuchFileException
 
 import firrtl.stage.TransformManager.TransformDependency
@@ -64,6 +64,39 @@ object FirrtlFileAnnotation extends HasShellOptions {
 
 }
 
+/** Read a directory of FIRRTL files or ProtoBufs
+  *  - set with `-I/--input-directory`
+  * @param file input filename
+  */
+case class FirrtlDirectoryAnnotation(dir: String) extends NoTargetAnnotation with CircuitOption {
+
+  def toCircuit(info: Parser.InfoMode): FirrtlCircuitAnnotation = {
+    val circuit =
+      try {
+        proto.FromProto.fromDirectory(dir)
+      } catch {
+        case a @ (_: FileNotFoundException | _: NoSuchFileException) =>
+          throw new OptionsException(s"Directory '$dir' not found! (Did you misspell it?)", a)
+      }
+    FirrtlCircuitAnnotation(circuit)
+  }
+
+}
+
+object FirrtlDirectoryAnnotation extends HasShellOptions {
+
+  val options = Seq(
+    new ShellOption[String](
+      longOption = "input-directory",
+      toAnnotationSeq = a => Seq(FirrtlDirectoryAnnotation(a)),
+      helpText = "A directory of FIRRTL files",
+      shortOption = Some("I"),
+      helpValueName = Some("<directory>")
+    )
+  )
+
+}
+
 /** An explicit output file the emitter will write to
   *   - set with `-o/--output-file`
   *  @param file output filename

src/main/scala/firrtl/stage/FirrtlCli.scala

@@ -13,6 +13,7 @@ trait FirrtlCli { this: Shell =>
   parser.note("FIRRTL Compiler Options")
   Seq(
     FirrtlFileAnnotation,
+    FirrtlDirectoryAnnotation,
     OutputFileAnnotation,
     InfoModeAnnotation,
     FirrtlSourceAnnotation,

src/main/scala/firrtl/stage/phases/Checks.scala

@@ -31,34 +31,39 @@ class Checks extends Phase {
     * @throws firrtl.options.OptionsException if any checks fail
     */
   def transform(annos: AnnotationSeq): AnnotationSeq = {
-    val inF, inS, eam, ec, outF, emitter, im, inC = collection.mutable.ListBuffer[Annotation]()
+    val inF, inS, inD, eam, ec, outF, emitter, im, inC = collection.mutable.ListBuffer[Annotation]()
     annos.foreach(_ match {
-      case a: FirrtlFileAnnotation     => a +=: inF
-      case a: FirrtlSourceAnnotation   => a +=: inS
-      case a: EmitAllModulesAnnotation => a +=: eam
-      case a: EmitCircuitAnnotation    => a +=: ec
-      case a: OutputFileAnnotation     => a +=: outF
-      case a: InfoModeAnnotation       => a +=: im
-      case a: FirrtlCircuitAnnotation  => a +=: inC
+      case a: FirrtlFileAnnotation      => a +=: inF
+      case a: FirrtlSourceAnnotation    => a +=: inS
+      case a: FirrtlDirectoryAnnotation => a +=: inD
+      case a: EmitAllModulesAnnotation  => a +=: eam
+      case a: EmitCircuitAnnotation     => a +=: ec
+      case a: OutputFileAnnotation      => a +=: outF
+      case a: InfoModeAnnotation        => a +=: im
+      case a: FirrtlCircuitAnnotation   => a +=: inC
       case a @ RunFirrtlTransformAnnotation(_: firrtl.Emitter) => a +=: emitter
       case _ =>
     })
 
     /* At this point, only a FIRRTL Circuit should exist */
-    if (inF.isEmpty && inS.isEmpty && inC.isEmpty) {
-      throw new OptionsException(s"""|Unable to determine FIRRTL source to read. None of the following were found:
-                                     |    - an input file:  -i, --input-file,    FirrtlFileAnnotation
-                                     |    - FIRRTL source:      --firrtl-source, FirrtlSourceAnnotation
-                                     |    - FIRRTL circuit:                      FirrtlCircuitAnnotation""".stripMargin)
+    if (inF.isEmpty && inS.isEmpty && inD.isEmpty && inC.isEmpty) {
+      throw new OptionsException(
+        s"""|Unable to determine FIRRTL source to read. None of the following were found:
+            |    - an input file:  -i, --input-file,      FirrtlFileAnnotation
+            |    - an input dir:   -I, --input-directory, FirrtlDirectoryAnnotation
+            |    - FIRRTL source:      --firrtl-source,   FirrtlSourceAnnotation
+            |    - FIRRTL circuit:                        FirrtlCircuitAnnotation""".stripMargin
+      )
     }
 
     /* Only one FIRRTL input can exist */
     if (inF.size + inS.size + inC.size > 1) {
       throw new OptionsException(
         s"""|Multiply defined input FIRRTL sources. More than one of the following was found:
-            |    - an input file (${inF.size} times):  -i, --input-file,    FirrtlFileAnnotation
-            |    - FIRRTL source (${inS.size} times):      --firrtl-source, FirrtlSourceAnnotation
-            |    - FIRRTL circuit (${inC.size} times):                      FirrtlCircuitAnnotation""".stripMargin
+            |    - an input file (${inF.size} times):  -i, --input-file,      FirrtlFileAnnotation
+            |    - an input dir (${inD.size} times):   -I, --input-directory, FirrtlDirectoryAnnotation
+            |    - FIRRTL source (${inS.size} times):      --firrtl-source,   FirrtlSourceAnnotation
+            |    - FIRRTL circuit (${inC.size} times):                        FirrtlCircuitAnnotation""".stripMargin
       )
     }