Skip to content

Commit

Permalink
Support serializing large (> 2 GiB) annotation files (#3905)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
jackkoenig authored Mar 5, 2024
1 parent 1c9c877 commit 736f1a7
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 5 deletions.
17 changes: 13 additions & 4 deletions firrtl/src/main/scala/firrtl/annotations/JsonProtocol.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
}
Expand Down
20 changes: 20 additions & 0 deletions firrtl/src/test/scala/firrtl/JsonProtocolSpec.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}

0 comments on commit 736f1a7

Please sign in to comment.