Let firrtl based applications run despite loading unknown annotations (#2387)

An application like barstools may contain a main that loads an annotations file containing
annotation classes that are not on it's classpath. This change allows unknown annotations
to be preserved by wrapping them in a UnrecognizedAnnotation. If annotations are then output
to a file, they will be unwrapped during serialization
This feature can be enabled via an AllowUnrecognizedAnnotations annotation

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

7 files changed, 125 insertions, 16 deletions

diff --git a/src/main/scala/firrtl/annotations/Annotation.scala b/src/main/scala/firrtl/annotations/Annotation.scala
index f4d6ee55..c6145c86 100644
--- a/src/main/scala/firrtl/annotations/Annotation.scala
+++ b/src/main/scala/firrtl/annotations/Annotation.scala
@@ -4,6 +4,7 @@ package firrtl
 package annotations
 
 import firrtl.options.StageUtils
+import org.json4s.JValue
 
 import scala.collection.Traversable
 
@@ -165,3 +166,5 @@ object Annotation
 case class DeletedAnnotation(xFormName: String, anno: Annotation) extends NoTargetAnnotation {
   override def serialize: String = s"""DELETED by $xFormName\n${anno.serialize}"""
 }
+
+case class UnrecognizedAnnotation(underlying: JValue) extends NoTargetAnnotation
diff --git a/src/main/scala/firrtl/annotations/AnnotationUtils.scala b/src/main/scala/firrtl/annotations/AnnotationUtils.scala
index 10894b26..826bd542 100644
--- a/src/main/scala/firrtl/annotations/AnnotationUtils.scala
+++ b/src/main/scala/firrtl/annotations/AnnotationUtils.scala
@@ -9,6 +9,7 @@ import firrtl.ir._
 
 case class InvalidAnnotationFileException(file: File, cause: FirrtlUserException = null)
     extends FirrtlUserException(s"$file", cause)
+case class UnrecogizedAnnotationsException(msg: String) extends FirrtlUserException(s"Unrecognized annotations $msg")
 case class InvalidAnnotationJSONException(msg: String) extends FirrtlUserException(msg)
 case class AnnotationFileNotFoundException(file: File)
     extends FirrtlUserException(
diff --git a/src/main/scala/firrtl/annotations/JsonProtocol.scala b/src/main/scala/firrtl/annotations/JsonProtocol.scala
index 3a25998d..6908a3a1 100644
--- a/src/main/scala/firrtl/annotations/JsonProtocol.scala
+++ b/src/main/scala/firrtl/annotations/JsonProtocol.scala
@@ -4,6 +4,8 @@ package firrtl
 package annotations
 
 import firrtl.ir._
+import firrtl.stage.AllowUnrecognizedAnnotations
+import logger.LazyLogging
 
 import scala.util.{Failure, Success, Try}
 
@@ -12,6 +14,8 @@ import org.json4s.native.JsonMethods._
 import org.json4s.native.Serialization
 import org.json4s.native.Serialization.{read, write, writePretty}
 
+import scala.collection.mutable
+
 trait HasSerializationHints {
   // For serialization of complicated constructor arguments, let the annotation
   // writer specify additional type hints for relevant classes that might be
@@ -22,7 +26,9 @@ trait HasSerializationHints {
 /** Wrapper [[Annotation]] for Annotations that cannot be serialized */
 case class UnserializeableAnnotation(error: String, content: String) extends NoTargetAnnotation
 
-object JsonProtocol {
+object JsonProtocol extends LazyLogging {
+  private val GetClassPattern = "[^']*'([^']+)'.*".r
+
   class TransformClassSerializer
       extends CustomSerializer[Class[_ <: Transform]](format =>
         (
@@ -207,6 +213,14 @@ object JsonProtocol {
         )
       )
 
+  class UnrecognizedAnnotationSerializer
+      extends CustomSerializer[JObject](format =>
+        (
+          { case JObject(s) => JObject(s) },
+          { case UnrecognizedAnnotation(underlying) => underlying }
+        )
+      )
+
   /** Construct Json formatter for annotations */
   def jsonFormat(tags: Seq[Class[_]]) = {
     Serialization.formats(FullTypeHints(tags.toList)).withTypeHintFieldName("class") +
@@ -217,7 +231,8 @@ object JsonProtocol {
       new LoadMemoryFileTypeSerializer + new IsModuleSerializer + new IsMemberSerializer +
       new CompleteTargetSerializer + new TypeSerializer + new ExpressionSerializer +
       new StatementSerializer + new PortSerializer + new DefModuleSerializer +
-      new CircuitSerializer + new InfoSerializer + new GroundTypeSerializer
+      new CircuitSerializer + new InfoSerializer + new GroundTypeSerializer +
+      new UnrecognizedAnnotationSerializer
   }
 
   /** Serialize annotations to a String for emission */
@@ -266,9 +281,17 @@ object JsonProtocol {
     writePretty(safeAnnos)
   }
 
-  def deserialize(in: JsonInput): Seq[Annotation] = deserializeTry(in).get
+  /** Deserialize JSON input into a Seq[Annotation]
+    *
+    * @param in JsonInput, can be file or string
+    * @param allowUnrecognizedAnnotations is set to true if command line contains flag to allow this behavior
+    * @return
+    */
+  def deserialize(in: JsonInput, allowUnrecognizedAnnotations: Boolean = false): Seq[Annotation] = {
+    deserializeTry(in, allowUnrecognizedAnnotations).get
+  }
 
-  def deserializeTry(in: JsonInput): Try[Seq[Annotation]] = Try({
+  def deserializeTry(in: JsonInput, allowUnrecognizedAnnotations: Boolean = false): Try[Seq[Annotation]] = Try {
     val parsed = parse(in)
     val annos = parsed match {
       case JArray(objs) => objs
@@ -277,6 +300,15 @@ object JsonProtocol {
           s"Annotations must be serialized as a JArray, got ${x.getClass.getName} instead!"
         )
     }
+
+    /* Tries to extract class name from the mapping exception */
+    def getAnnotationNameFromMappingException(mappingException: MappingException): String = {
+      mappingException.getMessage match {
+        case GetClassPattern(name) => name
+        case other                 => other
+      }
+    }
+
     // Recursively gather typeHints by pulling the "class" field from JObjects
     // Json4s should emit this as the first field in all serialized classes
     // Setting requireClassField mandates that all JObjects must provide a typeHint,
@@ -288,26 +320,83 @@ object JsonProtocol {
           throw new InvalidAnnotationJSONException(s"Expected field 'class' not found! $obj")
         case JObject(fields) => findTypeHints(fields.map(_._2))
         case JArray(arr)     => findTypeHints(arr)
-        case oJValue         => Seq()
+        case _               => Seq()
       })
       .distinct
 
+    // I don't much like this var here, but it has made it much simpler
+    // to maintain backward compatibility with the exception test structure
+    var classNotFoundBuildingLoaded = false
     val classes = findTypeHints(annos, true)
-    val loaded = classes.map(Class.forName(_))
+    val loaded = classes.flatMap { x =>
+      (try {
+        Some(Class.forName(x))
+      } catch {
+        case _: java.lang.ClassNotFoundException =>
+          classNotFoundBuildingLoaded = true
+          None
+      }): Option[Class[_]]
+    }
     implicit val formats = jsonFormat(loaded)
-    read[List[Annotation]](in)
-  }).recoverWith {
+    try {
+      read[List[Annotation]](in)
+    } catch {
+      case e: org.json4s.MappingException =>
+        // If we get here, the build `read` failed to process an annotation
+        // So we will map the annos one a time, wrapping the JSON of the unrecognized annotations
+        val exceptionList = new mutable.ArrayBuffer[String]()
+        val firrtlAnnos = annos.map { jsonAnno =>
+          try {
+            jsonAnno.extract[Annotation]
+          } catch {
+            case mappingException: org.json4s.MappingException =>
+              exceptionList += getAnnotationNameFromMappingException(mappingException)
+              UnrecognizedAnnotation(jsonAnno)
+          }
+        }
+
+        if (firrtlAnnos.contains(AllowUnrecognizedAnnotations) || allowUnrecognizedAnnotations) {
+          firrtlAnnos
+        } else {
+          logger.error(
+            "Annotation parsing found unrecognized annotations\n" +
+              "This error can be ignored with an AllowUnrecognizedAnnotationsAnnotation" +
+              " or command line flag --allow-unrecognized-annotations\n" +
+              exceptionList.mkString("\n")
+          )
+          if (classNotFoundBuildingLoaded) {
+            val distinctProblems = exceptionList.distinct
+            val problems = distinctProblems.take(10).mkString(", ")
+            val dots = if (distinctProblems.length > 10) {
+              ", ..."
+            } else {
+              ""
+            }
+            throw UnrecogizedAnnotationsException(s"($problems$dots)")
+          } else {
+            throw e
+          } // throw the mapping exception
+        }
+    }
+  }.recoverWith {
     // Translate some generic errors to specific ones
     case e: java.lang.ClassNotFoundException =>
-      Failure(new AnnotationClassNotFoundException(e.getMessage))
+      Failure(AnnotationClassNotFoundException(e.getMessage))
     // Eat the stack traces of json4s exceptions
     case e @ (_: org.json4s.ParserUtil.ParseException | _: org.json4s.MappingException) =>
-      Failure(new InvalidAnnotationJSONException(e.getMessage))
+      Failure(InvalidAnnotationJSONException(e.getMessage))
   }.recoverWith { // If the input is a file, wrap in InvalidAnnotationFileException
+    case e: UnrecogizedAnnotationsException =>
+      in match {
+        case FileInput(file) =>
+          Failure(InvalidAnnotationFileException(file, e))
+        case _ =>
+          Failure(e)
+      }
     case e: FirrtlUserException =>
       in match {
         case FileInput(file) =>
-          Failure(new InvalidAnnotationFileException(file, e))
+          Failure(InvalidAnnotationFileException(file, e))
         case _ => Failure(e)
       }
   }
diff --git a/src/main/scala/firrtl/options/phases/GetIncludes.scala b/src/main/scala/firrtl/options/phases/GetIncludes.scala
index 15692e79..d50b2c6f 100644
--- a/src/main/scala/firrtl/options/phases/GetIncludes.scala
+++ b/src/main/scala/firrtl/options/phases/GetIncludes.scala
@@ -6,9 +6,9 @@ import firrtl.AnnotationSeq
 import firrtl.annotations.{AnnotationFileNotFoundException, JsonProtocol}
 import firrtl.options.{InputAnnotationFileAnnotation, Phase, StageUtils}
 import firrtl.FileUtils
+import firrtl.stage.AllowUnrecognizedAnnotations
 
 import java.io.File
-
 import scala.collection.mutable
 import scala.util.{Failure, Try}
 
@@ -25,10 +25,13 @@ class GetIncludes extends Phase {
     * @param filename a JSON or YAML file of [[annotations.Annotation]]
     * @throws annotations.AnnotationFileNotFoundException if the file does not exist
     */
-  private def readAnnotationsFromFile(filename: String): AnnotationSeq = {
+  private def readAnnotationsFromFile(
+    filename:                     String,
+    allowUnrecognizedAnnotations: Boolean = false
+  ): AnnotationSeq = {
     val file = new File(filename).getCanonicalFile
     if (!file.exists) { throw new AnnotationFileNotFoundException(file) }
-    JsonProtocol.deserialize(file)
+    JsonProtocol.deserialize(file, allowUnrecognizedAnnotations)
   }
 
   /** Recursively read all [[Annotation]]s from any [[InputAnnotationFileAnnotation]]s while making sure that each file is
@@ -38,6 +41,7 @@ class GetIncludes extends Phase {
     * @return the original annotation sequence with any discovered annotations added
     */
   private def getIncludes(includeGuard: mutable.Set[String] = mutable.Set())(annos: AnnotationSeq): AnnotationSeq = {
+    val allowUnrecognizedAnnotations = annos.contains(AllowUnrecognizedAnnotations)
     annos.flatMap {
       case a @ InputAnnotationFileAnnotation(value) =>
         if (includeGuard.contains(value)) {
@@ -45,7 +49,7 @@ class GetIncludes extends Phase {
           None
         } else {
           includeGuard += value
-          getIncludes(includeGuard)(readAnnotationsFromFile(value))
+          getIncludes(includeGuard)(readAnnotationsFromFile(value, allowUnrecognizedAnnotations))
         }
       case x => Seq(x)
     }
diff --git a/src/main/scala/firrtl/stage/FirrtlAnnotations.scala b/src/main/scala/firrtl/stage/FirrtlAnnotations.scala
index 8e1e2001..6dbb0434 100644
--- a/src/main/scala/firrtl/stage/FirrtlAnnotations.scala
+++ b/src/main/scala/firrtl/stage/FirrtlAnnotations.scala
@@ -324,6 +324,16 @@ case object PrettyNoExprInlining extends NoTargetAnnotation with FirrtlOption wi
   )
 }
 
+case object AllowUnrecognizedAnnotations extends NoTargetAnnotation with FirrtlOption with HasShellOptions {
+  val options = Seq(
+    new ShellOption[Unit](
+      longOption = "allow-unrecognized-annotations",
+      toAnnotationSeq = _ => Seq(this),
+      helpText = "Allow annotation files to contain unrecognized annotations"
+    )
+  )
+}
+
 /** Turn off folding a specific primitive operand
   * @param op the op that should never be folded
   */
diff --git a/src/main/scala/firrtl/stage/FirrtlCli.scala b/src/main/scala/firrtl/stage/FirrtlCli.scala
index 2ad75574..a9c08627 100644
--- a/src/main/scala/firrtl/stage/FirrtlCli.scala
+++ b/src/main/scala/firrtl/stage/FirrtlCli.scala
@@ -27,7 +27,8 @@ trait FirrtlCli { this: Shell =>
     DisableFold,
     OptimizeForFPGA,
     CurrentFirrtlStateAnnotation,
-    CommonSubexpressionElimination
+    CommonSubexpressionElimination,
+    AllowUnrecognizedAnnotations
   )
     .map(_.addOptions(parser))
 
diff --git a/src/main/scala/firrtl/stage/package.scala b/src/main/scala/firrtl/stage/package.scala
index 92736963..b3fe2473 100644
--- a/src/main/scala/firrtl/stage/package.scala
+++ b/src/main/scala/firrtl/stage/package.scala
@@ -34,6 +34,7 @@ package object stage {
           case WarnNoScalaVersionDeprecation => c
           case PrettyNoExprInlining          => c
           case _: DisableFold => c
+          case AllowUnrecognizedAnnotations    => c
           case CurrentFirrtlStateAnnotation(a) => c
         }
       }