Support serializing large (> 2 GiB) annotation files (#3905) (#3907)

* Support serializing Annotations to a Writer Previous API serializing to a String is limited by JVM max size of 2 GiB Strings. * Serialize to Writer instead of String in WriteOutputAnnotations (cherry picked from commit 736f1a7) Co-authored-by: Jack Koenig <[email protected]>
chipsalliance · Mar 5, 2024 · 9769d54 · 9769d54
1 parent d71ea9a
commit 9769d54
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 5 deletions.
diff --git a/firrtl/src/main/scala/firrtl/annotations/JsonProtocol.scala b/firrtl/src/main/scala/firrtl/annotations/JsonProtocol.scala
@@ -15,6 +15,7 @@ import org.json4s.native.Serialization
 import org.json4s.native.Serialization.{read, write, writePretty}
 
 import scala.collection.mutable
+import java.io.{StringWriter, Writer}
 
 trait HasSerializationHints {
   // For serialization of complicated constructor arguments, let the annotation
@@ -167,13 +168,21 @@ object JsonProtocol extends LazyLogging {
       })
       .distinct
 
-  def serializeTry(annos: Seq[Annotation]): Try[String] = {
-    val tags = getTags(annos)
+  def serializeTry(annos: Seq[Annotation]): Try[String] = serializeTry(annos, new StringWriter).map(_.toString)
+
+  /** Serialize annotations to a [[java.io.Writer]]
+    *
+    * @param annos Annotations to serialize
+    * @param out Writer to which the serialized annotations will be written
+    * @return
+    */
+  def serializeTry[W <: Writer](annos: Iterable[Annotation], out: W): Try[W] = {
+    val tags = getTags(annos.toSeq)
 
     implicit val formats = jsonFormat(tags)
-    Try(writePretty(annos)).recoverWith {
+    Try(writePretty(annos, out)).recoverWith {
       case e: org.json4s.MappingException =>
-        val badAnnos = findUnserializeableAnnos(annos)
+        val badAnnos = findUnserializeableAnnos(annos.toSeq)
         Failure(if (badAnnos.isEmpty) e else UnserializableAnnotationException(badAnnos))
     }
   }

diff --git a/firrtl/src/main/scala/firrtl/options/phases/WriteOutputAnnotations.scala b/firrtl/src/main/scala/firrtl/options/phases/WriteOutputAnnotations.scala
@@ -126,7 +126,7 @@ class WriteOutputAnnotations extends Phase {
           case None =>
           case Some(file) =>
             val pw = new PrintWriter(sopts.getBuildFileName(file, Some(".anno.json")))
-            pw.write(JsonProtocol.serialize(serializable))
+            JsonProtocol.serializeTry(serializable, pw).get // .get to throw any exceptions
             pw.close()
         }
     }

diff --git a/firrtl/src/test/scala/firrtl/JsonProtocolSpec.scala b/firrtl/src/test/scala/firrtl/JsonProtocolSpec.scala
@@ -80,4 +80,24 @@ class JsonProtocolSpec extends AnyFlatSpec {
     val deserAnno = JsonProtocol.deserialize(serializedAnno).head
     assert(anno == deserAnno)
   }
+
+  "JsonProtocol" should "support serializing directly to a Java Writer" in {
+    val anno = SimpleAnnotation("hello")
+    class NaiveWriter extends java.io.Writer {
+      private var contents: String = ""
+      def value:            String = contents
+      def close():          Unit = contents = ""
+      def flush():          Unit = contents = ""
+      def write(cbuff: Array[Char], off: Int, len: Int): Unit = {
+        for (i <- off until off + len) {
+          contents += cbuff(i)
+        }
+      }
+    }
+    val w = new NaiveWriter
+    JsonProtocol.serializeTry(Seq(anno), w)
+    val ser1 = w.value
+    val ser2 = JsonProtocol.serialize(Seq(anno))
+    assert(ser1 == ser2)
+  }
 }