From c2d72fd8a4f3558094094eed2f4a12f85d080c41 Mon Sep 17 00:00:00 2001
From: Jack Koenig <koenig@sifive.com>
Date: Fri, 14 May 2021 11:12:19 -0700
Subject: [PATCH] Add JsonProtocol.serializeRecover (#2227)

This function will safely wrap any unserializeable annotations in
UnserializeableAnnotations so that they can be safely serialized to JSON
for logging.
---
 src/main/scala/firrtl/Compiler.scala          | 12 +-------
 .../firrtl/annotations/JsonProtocol.scala     | 29 +++++++++++++++++--
 .../annotationTests/JsonProtocolSpec.scala    | 12 +++++++-
 3 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/src/main/scala/firrtl/Compiler.scala b/src/main/scala/firrtl/Compiler.scala
index 2998af3c6e..7354b8ee8c 100644
--- a/src/main/scala/firrtl/Compiler.scala
+++ b/src/main/scala/firrtl/Compiler.scala
@@ -214,17 +214,7 @@ private[firrtl] object Transform {
     val remappedAnnotations = propagateAnnotations(name, logger, before.annotations, after.annotations, after.renames)
 
     logger.trace(s"Annotations:")
-    logger.trace {
-      JsonProtocol
-        .serializeTry(remappedAnnotations)
-        .recoverWith {
-          case NonFatal(e) =>
-            val msg = s"Exception thrown during Annotation serialization:\n  " +
-              e.toString.replaceAll("\n", "\n  ")
-            Try(msg)
-        }
-        .get
-    }
+    logger.trace(JsonProtocol.serializeRecover(remappedAnnotations))
 
     logger.trace(s"Circuit:\n${after.circuit.serialize}")
 
diff --git a/src/main/scala/firrtl/annotations/JsonProtocol.scala b/src/main/scala/firrtl/annotations/JsonProtocol.scala
index c284c788f4..3a25998dd4 100644
--- a/src/main/scala/firrtl/annotations/JsonProtocol.scala
+++ b/src/main/scala/firrtl/annotations/JsonProtocol.scala
@@ -5,7 +5,7 @@ package annotations
 
 import firrtl.ir._
 
-import scala.util.{Failure, Try}
+import scala.util.{Failure, Success, Try}
 
 import org.json4s._
 import org.json4s.native.JsonMethods._
@@ -19,6 +19,9 @@ trait HasSerializationHints {
   def typeHints: Seq[Class[_]]
 }
 
+/** Wrapper [[Annotation]] for Annotations that cannot be serialized */
+case class UnserializeableAnnotation(error: String, content: String) extends NoTargetAnnotation
+
 object JsonProtocol {
   class TransformClassSerializer
       extends CustomSerializer[Class[_ <: Transform]](format =>
@@ -227,14 +230,17 @@ object JsonProtocol {
   ): Seq[(Annotation, Throwable)] =
     annos.map(a => a -> Try(write(a))).collect { case (a, Failure(e)) => (a, e) }
 
-  def serializeTry(annos: Seq[Annotation]): Try[String] = {
-    val tags = annos
+  private def getTags(annos: Seq[Annotation]): Seq[Class[_]] =
+    annos
       .flatMap({
         case anno: HasSerializationHints => anno.getClass +: anno.typeHints
         case anno => Seq(anno.getClass)
       })
       .distinct
 
+  def serializeTry(annos: Seq[Annotation]): Try[String] = {
+    val tags = getTags(annos)
+
     implicit val formats = jsonFormat(tags)
     Try(writePretty(annos)).recoverWith {
       case e: org.json4s.MappingException =>
@@ -243,6 +249,23 @@ object JsonProtocol {
     }
   }
 
+  /** Serialize annotations to JSON while wrapping unserializeable ones with [[UnserializeableAnnotation]]
+    *
+    * @note this is slower than standard serialization
+    */
+  def serializeRecover(annos: Seq[Annotation]): String = {
+    val tags = classOf[UnserializeableAnnotation] +: getTags(annos)
+    implicit val formats = jsonFormat(tags)
+
+    val safeAnnos = annos.map { anno =>
+      Try(write(anno)) match {
+        case Success(_) => anno
+        case Failure(e) => UnserializeableAnnotation(e.getMessage, anno.toString)
+      }
+    }
+    writePretty(safeAnnos)
+  }
+
   def deserialize(in: JsonInput): Seq[Annotation] = deserializeTry(in).get
 
   def deserializeTry(in: JsonInput): Try[Seq[Annotation]] = Try({
diff --git a/src/test/scala/firrtlTests/annotationTests/JsonProtocolSpec.scala b/src/test/scala/firrtlTests/annotationTests/JsonProtocolSpec.scala
index 6dd018ebde..b2d029f579 100644
--- a/src/test/scala/firrtlTests/annotationTests/JsonProtocolSpec.scala
+++ b/src/test/scala/firrtlTests/annotationTests/JsonProtocolSpec.scala
@@ -3,9 +3,10 @@
 package firrtlTests.annotationTests
 
 import firrtl._
-import firrtl.annotations.{JsonProtocol, NoTargetAnnotation, UnserializableAnnotationException}
+import firrtl.annotations._
 import firrtl.ir._
 import firrtl.options.Dependency
+import firrtl.transforms.DontTouchAnnotation
 import scala.util.Failure
 import _root_.logger.{LogLevel, LogLevelAnnotation, Logger}
 import org.scalatest.flatspec.AnyFlatSpec
@@ -77,4 +78,13 @@ class JsonProtocolSpec extends AnyFlatSpec with Matchers {
         e.getMessage should include("Classes defined in method bodies are not supported")
     }
   }
+  "JsonProtocol.serializeRecover" should "emit even annotations that cannot be serialized" in {
+    case class MyAnno(x: Int) extends NoTargetAnnotation
+    val target = CircuitTarget("Top").module("Foo").ref("x")
+    val annos = MyAnno(3) :: DontTouchAnnotation(target) :: Nil
+    val res = JsonProtocol.serializeRecover(annos)
+    res should include(""""class":"firrtl.annotations.UnserializeableAnnotation",""")
+    res should include(""""error":"Classes defined in method bodies are not supported.",""")
+    res should include(""""content":"MyAnno(3)"""")
+  }
 }