Multi protobuf module emission and consumption (#2344)

* Add compiler option (`-p`) to emit individual module protobufs
* Implement multi module combination when reading directory of protobufs

Co-authored-by: Jack Koenig <koenig@sifive.com>

10 files changed, 443 insertions, 36 deletions

diff --git a/src/main/scala/firrtl/Emitter.scala b/src/main/scala/firrtl/Emitter.scala
index 589acc70..0814c268 100644
--- a/src/main/scala/firrtl/Emitter.scala
+++ b/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-modules-protobuf",
+      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>")
     )
   )
 
diff --git a/src/main/scala/firrtl/Utils.scala b/src/main/scala/firrtl/Utils.scala
index e2bb06ff..51748b68 100644
--- a/src/main/scala/firrtl/Utils.scala
+++ b/src/main/scala/firrtl/Utils.scala
@@ -978,6 +978,65 @@ object Utils extends LazyLogging {
     map.view.map({ case (k, vs) => k -> vs.toList }).toList
   }
 
+  // For a given module, returns a Seq of all instantiated modules inside of it
+  private[firrtl] def collectInstantiatedModules(mod: Module, map: Map[String, DefModule]): Seq[DefModule] = {
+    // Use list instead of set to maintain order
+    val modules = mutable.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
+  }
+
+  /** Checks if two circuits are equal regardless of their ordering of module definitions */
+  def orderAgnosticEquality(a: Circuit, b: Circuit): Boolean =
+    a.copy(modules = a.modules.sortBy(_.name)) == b.copy(modules = b.modules.sortBy(_.name))
+
+  /** 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 module with no ExtModules, Right means child modules or lone ExtModules
+      val module: Option[Module] = {
+        val found: Seq[Module] = modules.collect { case m: Module => m }
+        assert(
+          found.size <= 1,
+          s"Module definitions should have unique names, found ${found.size} definitions named ${found.head.name}"
+        )
+        found.headOption
+      }
+      val extModules: Seq[ExtModule] = modules.collect { case e: ExtModule => e }.distinct
+
+      // If the module is a lone module (no extmodule references in any other file)
+      if (extModules.isEmpty && !module.isEmpty)
+        Seq(Left(module.get))
+      // If a module has extmodules, but no other file contains the implementation
+      else if (!extModules.isEmpty && module.isEmpty)
+        extModules.map(Right(_))
+      // Otherwise there is a module implementation with extmodule references
+      else
+        Seq(Right(module.get))
+    }
+
+    // 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, s"There should only be 1 top module, got: ${found.map(_.name).mkString(", ")}")
+      found.head
+    }
+    val res = deduped.collect { case Right(m) => m }
+    ir.Circuit(NoInfo, top +: res.toSeq, top.name)
+  }
+
   object True {
     private val _True = UIntLiteral(1, IntWidth(1))
 
diff --git a/src/main/scala/firrtl/backends/firrtl/FirrtlEmitter.scala b/src/main/scala/firrtl/backends/firrtl/FirrtlEmitter.scala
index 56b63d75..69f3da00 100644
--- a/src/main/scala/firrtl/backends/firrtl/FirrtlEmitter.scala
+++ b/src/main/scala/firrtl/backends/firrtl/FirrtlEmitter.scala
@@ -18,18 +18,6 @@ sealed abstract class FirrtlEmitter(form: Seq[TransformDependency], val outputSu
   override def invalidates(a: Transform) = false
 
   private def emitAllModules(circuit: Circuit): Seq[EmittedFirrtlModule] = {
-    // 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 = mutable.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 {
diff --git a/src/main/scala/firrtl/backends/proto/ProtoBufEmitter.scala b/src/main/scala/firrtl/backends/proto/ProtoBufEmitter.scala
index 31edb6b8..c617ea27 100644
--- a/src/main/scala/firrtl/backends/proto/ProtoBufEmitter.scala
+++ b/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.{collectInstantiatedModules, throwInternalError}
 
 /** This object defines Annotations that are used by Protocol Buffer emission.
   */
@@ -59,10 +62,35 @@ 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] = {
+    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())
diff --git a/src/main/scala/firrtl/proto/FromProto.scala b/src/main/scala/firrtl/proto/FromProto.scala
index ed641eb2..91b3f872 100644
--- a/src/main/scala/firrtl/proto/FromProto.scala
+++ b/src/main/scala/firrtl/proto/FromProto.scala
@@ -9,6 +9,11 @@ import collection.JavaConverters._
 import FirrtlProtos._
 import com.google.protobuf.CodedInputStream
 import Firrtl.Statement.{Formal, ReadUnderWrite}
+import firrtl.ir.DefModule
+import Utils.combine
+import java.io.FileNotFoundException
+import firrtl.options.OptionsException
+import java.nio.file.NotDirectoryException
 
 object FromProto {
 
@@ -33,6 +38,26 @@ object FromProto {
     proto.FromProto.convert(pb)
   }
 
+  /** Deserialize all the ProtoBuf representations of [[ir.Circuit]] in @dir
+    *
+    * @param dir directory containing ProtoBuf representation(s)
+    * @return Deserialized FIRRTL Circuit
+    * @throws java.io.FileNotFoundException if dir does not exist
+    * @throws java.nio.file.NotDirectoryException if dir exists but is not a directory
+    */
+  def fromDirectory(dir: String): ir.Circuit = {
+    val d = new File(dir)
+    if (!d.exists) {
+      throw new FileNotFoundException(s"Specified directory '$d' does not exist!")
+    }
+    if (!d.isDirectory) {
+      throw new NotDirectoryException(s"'$d' is not a directory!")
+    }
+
+    val fileList = d.listFiles.filter(_.isFile).toList
+    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
diff --git a/src/main/scala/firrtl/stage/FirrtlAnnotations.scala b/src/main/scala/firrtl/stage/FirrtlAnnotations.scala
index 9e6fefca..8e1e2001 100644
--- a/src/main/scala/firrtl/stage/FirrtlAnnotations.scala
+++ b/src/main/scala/firrtl/stage/FirrtlAnnotations.scala
@@ -6,8 +6,8 @@ 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.nio.file.NoSuchFileException
+import java.io.{File, FileNotFoundException}
+import java.nio.file.{NoSuchFileException, NotDirectoryException}
 
 import firrtl.stage.TransformManager.TransformDependency
 
@@ -64,6 +64,43 @@ object FirrtlFileAnnotation extends HasShellOptions {
 
 }
 
+/** Read a directory of ProtoBufs
+  *  - set with `-I/--input-directory`
+  *
+  * TODO: Does not currently support FIRRTL files.
+  * @param dir input directory name
+  */
+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)
+        case _: NotDirectoryException =>
+          throw new OptionsException(s"Directory '$dir' is not a directory")
+      }
+    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
diff --git a/src/main/scala/firrtl/stage/FirrtlCli.scala b/src/main/scala/firrtl/stage/FirrtlCli.scala
index abba0c66..2ad75574 100644
--- a/src/main/scala/firrtl/stage/FirrtlCli.scala
+++ b/src/main/scala/firrtl/stage/FirrtlCli.scala
@@ -14,6 +14,7 @@ trait FirrtlCli { this: Shell =>
   parser.note("FIRRTL Compiler Options")
   Seq(
     FirrtlFileAnnotation,
+    FirrtlDirectoryAnnotation,
     OutputFileAnnotation,
     InfoModeAnnotation,
     FirrtlSourceAnnotation,
diff --git a/src/main/scala/firrtl/stage/phases/Checks.scala b/src/main/scala/firrtl/stage/phases/Checks.scala
index e9065393..edfb7d03 100644
--- a/src/main/scala/firrtl/stage/phases/Checks.scala
+++ b/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
       )
     }
 
diff --git a/src/test/scala/firrtlTests/UtilsSpec.scala b/src/test/scala/firrtlTests/UtilsSpec.scala
index 99f1ffd0..1048b370 100644
--- a/src/test/scala/firrtlTests/UtilsSpec.scala
+++ b/src/test/scala/firrtlTests/UtilsSpec.scala
@@ -46,4 +46,174 @@ class UtilsSpec extends AnyFlatSpec {
 
     (Utils.expandRef(wr)) should be(expected)
   }
+
+  def combineTest(circuits: Seq[String], expected: String) = {
+    (Utils.orderAgnosticEquality(Utils.combine(circuits.map(c => Parser.parse(c))), Parser.parse(expected))) should be(
+      true
+    )
+  }
+
+  "combine" should "merge multiple module circuits" in {
+    val input = Seq(
+      """|circuit Top:
+         |  extmodule External:
+         |    output foo: UInt<32>
+         |    defname = External
+         |
+         |  extmodule Child1:
+         |    output foo: UInt<32>
+         |    defname = Child1
+         |
+         |  extmodule Child2:
+         |    output foo: UInt<32>
+         |    defname = Child2
+         |
+         |  module Top:
+         |    output foo: UInt<32>
+         |    inst c1 of Child1
+         |    inst c2 of Child2
+         |    inst e of External
+         |    foo <= tail(add(add(c1.foo, c2.bar), e.foo), 1)
+         |""".stripMargin,
+      """|circuit Child1:
+         |  extmodule External:
+         |    output foo: UInt<32>
+         |    defname = External
+         |
+         |  module Child1:
+         |    output foo: UInt<32>
+         |    inst e of External
+         |    foo <= e.foo
+         |""".stripMargin,
+      """|circuit Child2:
+         |  extmodule External:
+         |    output foo: UInt<32>
+         |    defname = External
+         |
+         |  module Child2:
+         |    output bar: UInt<32>
+         |    inst e of External
+         |    bar <= e.foo
+         |""".stripMargin
+    )
+
+    val output =
+      """|circuit Top:
+         |  module Top:
+         |    output foo: UInt<32>
+         |    inst c1 of Child1
+         |    inst c2 of Child2
+         |    inst e of External
+         |    foo <= tail(add(add(c1.foo, c2.bar), e.foo), 1)
+         |
+         |  module Child1:
+         |    output foo: UInt<32>
+         |    inst e of External
+         |    foo <= e.foo
+         |
+         |  module Child2:
+         |    output bar: UInt<32>
+         |    inst e of External
+         |    bar <= e.foo
+         |
+         |  extmodule External:
+         |    output foo: UInt<32>
+         |    defname = External
+         |""".stripMargin
+
+    combineTest(input, output)
+  }
+
+  "combine" should "dedup ExtModules if an implementation exists" in {
+    val input = Seq(
+      """|circuit Top:
+         |  extmodule Child:
+         |    output foo: UInt<32>
+         |    defname = Child
+         |
+         |  module Top:
+         |    output foo: UInt<32>
+         |    inst c of Child
+         |    foo <= c.foo
+         |""".stripMargin,
+      """|circuit Child:
+         |  module Child:
+         |    output foo: UInt<32>
+         |
+         |    skip
+         |""".stripMargin
+    )
+
+    val output =
+      """|circuit Top:
+         |  module Top:
+         |    output foo: UInt<32>
+         |    inst c of Child
+         |    foo <= c.foo
+         |
+         |  module Child:
+         |    output foo: UInt<32>
+         |
+         |    skip
+         |""".stripMargin
+
+    combineTest(input, output)
+  }
+
+  "combine" should "support lone ExtModules" in {
+    val input = Seq(
+      """|circuit Top:
+         |  extmodule External:
+         |    output foo: UInt<32>
+         |    defname = External
+         |
+         |  module Top:
+         |    output foo: UInt<32>
+         |    inst e of External
+         |    foo <= e.foo
+         |""".stripMargin
+    )
+
+    val output =
+      """|circuit Top:
+         |  extmodule External:
+         |    output foo: UInt<32>
+         |    defname = External
+         |
+         |  module Top:
+         |    output foo: UInt<32>
+         |    inst e of External
+         |    foo <= e.foo
+         |""".stripMargin
+
+    combineTest(input, output)
+  }
+
+  "combine" should "fail with multiple lone Modules" in {
+    val input = Seq(
+      """|circuit Top:
+         |  extmodule External:
+         |    output foo: UInt<32>
+         |    defname = External
+         |
+         |  module Top:
+         |    output foo: UInt<32>
+         |    inst e of External
+         |    foo <= e.foo
+         |""".stripMargin,
+      """|circuit Top2:
+         |  extmodule External:
+         |    output foo: UInt<32>
+         |    defname = External
+         |
+         |  module Top2:
+         |    output bar: UInt<32>
+         |    inst e of External
+         |    bar <= e.foo
+         |""".stripMargin
+    )
+
+    a[java.lang.AssertionError] shouldBe thrownBy { combineTest(input, "") }
+  }
+
 }
diff --git a/src/test/scala/firrtlTests/stage/FirrtlMainSpec.scala b/src/test/scala/firrtlTests/stage/FirrtlMainSpec.scala
index d51fc164..30e03b3c 100644
--- a/src/test/scala/firrtlTests/stage/FirrtlMainSpec.scala
+++ b/src/test/scala/firrtlTests/stage/FirrtlMainSpec.scala
@@ -246,6 +246,16 @@ class FirrtlMainSpec
         args = Array("-X", "sverilog", "-e", "sverilog"),
         files = Seq("Top.sv", "Child.sv")
       ),
+      /* Test all one protobuf per module emitters */
+      FirrtlMainTest(
+        args = Array("-X", "none", "--emit-modules-protobuf", "chirrtl"),
+        files = Seq("Top.pb", "Child.pb")
+      ),
+      FirrtlMainTest(args = Array("-X", "none", "-p", "mhigh"), files = Seq("Top.mhi.pb", "Child.mhi.pb")),
+      FirrtlMainTest(args = Array("-X", "none", "-p", "high"), files = Seq("Top.hi.pb", "Child.hi.pb")),
+      FirrtlMainTest(args = Array("-X", "none", "-p", "middle"), files = Seq("Top.mid.pb", "Child.mid.pb")),
+      FirrtlMainTest(args = Array("-X", "none", "-p", "low"), files = Seq("Top.lo.pb", "Child.lo.pb")),
+      FirrtlMainTest(args = Array("-X", "none", "-p", "low-opt"), files = Seq("Top.lo.pb", "Child.lo.pb")),
       /* Test mixing of -E with -e */
       FirrtlMainTest(
         args = Array("-X", "middle", "-E", "high", "-e", "middle"),
@@ -324,6 +334,47 @@ class FirrtlMainSpec
       new File(td.buildDir + "/Foo.hi.fir") should (exist)
     }
 
+    Scenario("User compiles to multiple Protocol Buffers") {
+      val f = new FirrtlMainFixture
+      val td = new TargetDirectoryFixture("multi-protobuf")
+      val c = new SimpleFirrtlCircuitFixture
+      val protobufs = Seq("Top.pb", "Child.pb")
+
+      And("some input multi-module FIRRTL IR")
+      val inputFile: Array[String] = {
+        val in = new File(td.dir, c.main)
+        val pw = new PrintWriter(in)
+        pw.write(c.input)
+        pw.close()
+        Array("-i", in.toString)
+      }
+
+      When("the user tries to emit a circuit to multiple Protocol Buffer files in the target directory")
+      f.stage.main(
+        inputFile ++ Array("-X", "none", "-p", "chirrtl", "-td", td.buildDir.toString)
+      )
+
+      protobufs.foreach { f =>
+        Then(s"file '$f' should be emitted")
+        val out = new File(td.buildDir + s"/$f")
+        out should (exist)
+      }
+
+      When("the user compiles the Protobufs to a single FIRRTL IR")
+      f.stage.main(
+        Array("-I", td.buildDir.toString, "-X", "none", "-E", "chirrtl", "-td", td.buildDir.toString, "-o", "Foo")
+      )
+
+      Then("one single FIRRTL file should be emitted")
+      val outFile = new File(td.buildDir + "/Foo.fir")
+      outFile should (exist)
+      And("it should be the same as using FIRRTL input")
+      firrtl.Utils.orderAgnosticEquality(
+        firrtl.Parser.parse(c.input),
+        firrtl.Parser.parseFile(td.buildDir + "/Foo.fir", firrtl.Parser.IgnoreInfo)
+      ) should be(true)
+    }
+
   }
 
   info("As a FIRRTL command line user")
@@ -381,6 +432,12 @@ class FirrtlMainSpec
         circuit = None,
         stdout = Some("Unknown compiler name 'Verilog'! (Did you misspell it?)"),
         result = 1
+      ),
+      FirrtlMainTest(
+        args = Array("-I", "test_run_dir/I-DO-NOT-EXIST"),
+        circuit = None,
+        stdout = Some("Directory 'test_run_dir/I-DO-NOT-EXIST' not found!"),
+        result = 1
       )
     )
       .foreach(runStageExpectFiles)