Merge pull request #1639 from freechipsproject/improve-anno-logging

Improve annotation logging

3 files changed, 117 insertions, 129 deletions

diff --git a/src/main/scala/firrtl/Compiler.scala b/src/main/scala/firrtl/Compiler.scala
index 06bd512e..6f921189 100644
--- a/src/main/scala/firrtl/Compiler.scala
+++ b/src/main/scala/firrtl/Compiler.scala
@@ -5,8 +5,9 @@ package firrtl
 import logger._
 import java.io.Writer
 
-
 import scala.collection.mutable
+import scala.util.Try
+import scala.util.control.NonFatal
 
 import firrtl.annotations._
 import firrtl.ir.Circuit
@@ -192,6 +193,88 @@ final case object UnknownForm extends CircuitForm(-1) {
 }
 // scalastyle:on magic.number
 
+// Internal utilities to keep code DRY, not a clean interface
+private[firrtl] object Transform {
+
+  // Run transform with logging
+  def runTransform(name: String, mk: => CircuitState, logger: Logger): CircuitState = {
+    logger.info(s"======== Starting Transform $name ========")
+
+    val (timeMillis, result) = Utils.time(mk)
+
+    logger.info(s"""----------------------------${"-" * name.size}---------\n""")
+    logger.info(f"Time: $timeMillis%.1f ms")
+
+    result
+  }
+
+  def remapAnnotations(name: String, before: CircuitState, after: CircuitState, logger: Logger): CircuitState = {
+    val remappedAnnotations = propagateAnnotations(name, logger, before.annotations, after.annotations, after.renames)
+
+    logger.info(s"Form: ${after.form}")
+    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(s"Circuit:\n${after.circuit.serialize}")
+    logger.info(s"======== Finished Transform $name ========\n")
+
+    CircuitState(after.circuit, after.form, remappedAnnotations, None)
+  }
+
+  /** Propagate annotations and update their names.
+    *
+    * @param inAnno input AnnotationSeq
+    * @param resAnno result AnnotationSeq
+    * @param renameOpt result RenameMap
+    * @return the updated annotations
+    */
+  def propagateAnnotations(
+      name: String,
+      logger: Logger,
+      inAnno: AnnotationSeq,
+      resAnno: AnnotationSeq,
+      renameOpt: Option[RenameMap]): AnnotationSeq = {
+    val newAnnotations = {
+      val inSet = mutable.LinkedHashSet() ++ inAnno
+      val resSet = mutable.LinkedHashSet() ++ resAnno
+      val deleted = (inSet -- resSet).map {
+        case DeletedAnnotation(xFormName, delAnno) => DeletedAnnotation(s"$xFormName+$name", delAnno)
+        case anno => DeletedAnnotation(name, anno)
+      }
+      val created = resSet -- inSet
+      val unchanged = resSet & inSet
+      (deleted ++ created ++ unchanged)
+    }
+
+    // For each annotation, rename all annotations.
+    val renames = renameOpt.getOrElse(RenameMap())
+    val remapped2original = mutable.LinkedHashMap[Annotation, mutable.LinkedHashSet[Annotation]]()
+    val keysOfNote = mutable.LinkedHashSet[Annotation]()
+    val finalAnnotations = newAnnotations.flatMap { anno =>
+      val remappedAnnos = anno.update(renames)
+      remappedAnnos.foreach { remapped =>
+        val set = remapped2original.getOrElseUpdate(remapped, mutable.LinkedHashSet.empty[Annotation])
+        set += anno
+        if(set.size > 1) keysOfNote += remapped
+      }
+      remappedAnnos
+    }.toSeq
+    keysOfNote.foreach { key =>
+      logger.debug(s"""The following original annotations are renamed to the same new annotation.""")
+      logger.debug(s"""Original Annotations:\n  ${remapped2original(key).mkString("\n  ")}""")
+      logger.debug(s"""New Annotation:\n  $key""")
+    }
+    finalAnnotations
+  }
+}
+
 /** The basic unit of operating on a Firrtl AST */
 trait Transform extends TransformLike[CircuitState] with DependencyAPI[Transform] {
 
@@ -299,66 +382,10 @@ trait Transform extends TransformLike[CircuitState] with DependencyAPI[Transform
     * @return A transformed Firrtl AST
     */
   final def runTransform(state: CircuitState): CircuitState = {
-    logger.info(s"======== Starting Transform $name ========")
-
-    val (timeMillis, result) = Utils.time { execute(prepare(state)) }
-
-    logger.info(s"""----------------------------${"-" * name.size}---------\n""")
-    logger.info(f"Time: $timeMillis%.1f ms")
-
-    val remappedAnnotations = propagateAnnotations(state.annotations, result.annotations, result.renames)
-
-    logger.info(s"Form: ${result.form}")
-    logger.trace(s"Annotations:")
-    logger.trace(JsonProtocol.serialize(remappedAnnotations))
-    logger.trace(s"Circuit:\n${result.circuit.serialize}")
-    logger.info(s"======== Finished Transform $name ========\n")
-    CircuitState(result.circuit, result.form, remappedAnnotations, None)
+    val result = Transform.runTransform(name, execute(prepare(state)), logger)
+    Transform.remapAnnotations(name, state, result, logger)
   }
 
-  /** Propagate annotations and update their names.
-    *
-    * @param inAnno input AnnotationSeq
-    * @param resAnno result AnnotationSeq
-    * @param renameOpt result RenameMap
-    * @return the updated annotations
-    */
-  final private def propagateAnnotations(
-      inAnno: AnnotationSeq,
-      resAnno: AnnotationSeq,
-      renameOpt: Option[RenameMap]): AnnotationSeq = {
-    val newAnnotations = {
-      val inSet = mutable.LinkedHashSet() ++ inAnno
-      val resSet = mutable.LinkedHashSet() ++ resAnno
-      val deleted = (inSet -- resSet).map {
-        case DeletedAnnotation(xFormName, delAnno) => DeletedAnnotation(s"$xFormName+$name", delAnno)
-        case anno => DeletedAnnotation(name, anno)
-      }
-      val created = resSet -- inSet
-      val unchanged = resSet & inSet
-      (deleted ++ created ++ unchanged)
-    }
-
-    // For each annotation, rename all annotations.
-    val renames = renameOpt.getOrElse(RenameMap())
-    val remapped2original = mutable.LinkedHashMap[Annotation, mutable.LinkedHashSet[Annotation]]()
-    val keysOfNote = mutable.LinkedHashSet[Annotation]()
-    val finalAnnotations = newAnnotations.flatMap { anno =>
-      val remappedAnnos = anno.update(renames)
-      remappedAnnos.foreach { remapped =>
-        val set = remapped2original.getOrElseUpdate(remapped, mutable.LinkedHashSet.empty[Annotation])
-        set += anno
-        if(set.size > 1) keysOfNote += remapped
-      }
-      remappedAnnos
-    }.toSeq
-    keysOfNote.foreach { key =>
-      logger.debug(s"""The following original annotations are renamed to the same new annotation.""")
-      logger.debug(s"""Original Annotations:\n  ${remapped2original(key).mkString("\n  ")}""")
-      logger.debug(s"""New Annotation:\n  $key""")
-    }
-    finalAnnotations
-  }
 }
 
 trait SeqTransformBased {
diff --git a/src/main/scala/firrtl/stage/transforms/UpdateAnnotations.scala b/src/main/scala/firrtl/stage/transforms/UpdateAnnotations.scala
index e9879979..cc0fbc6f 100644
--- a/src/main/scala/firrtl/stage/transforms/UpdateAnnotations.scala
+++ b/src/main/scala/firrtl/stage/transforms/UpdateAnnotations.scala
@@ -2,12 +2,9 @@
 
 package firrtl.stage.transforms
 
-import firrtl.{AnnotationSeq, CircuitState, RenameMap, Transform, Utils}
-import firrtl.annotations.{Annotation, DeletedAnnotation, JsonProtocol}
+import firrtl.{CircuitState, Transform}
 import firrtl.options.Translator
 
-import scala.collection.mutable
-
 class UpdateAnnotations(val underlying: Transform) extends Transform with WrappedTransform
     with Translator[CircuitState, (CircuitState, CircuitState)] {
 
@@ -16,74 +13,13 @@ class UpdateAnnotations(val underlying: Transform) extends Transform with Wrappe
   def aToB(a: CircuitState): (CircuitState, CircuitState) = (a, a)
 
   def bToA(b: (CircuitState, CircuitState)): CircuitState = {
-    val (state, result) = (b._1, b._2)
-
-    val remappedAnnotations = propagateAnnotations(state.annotations, result.annotations, result.renames)
-
-    logger.info(s"Form: ${result.form}")
-    logger.trace(s"Annotations:")
-    logger.trace(JsonProtocol.serialize(remappedAnnotations))
-
-    logger.trace(s"Circuit:\n${result.circuit.serialize}")
-    logger.info(s"======== Finished Transform $name ========\n")
-
-    CircuitState(result.circuit, result.form, remappedAnnotations, None)
+    Transform.remapAnnotations(name, b._1, b._2, logger)
   }
 
   def internalTransform(b: (CircuitState, CircuitState)): (CircuitState, CircuitState) = {
-    logger.info(s"======== Starting Transform $name ========")
-
-    val (timeMillis, result) = Utils.time { underlying.transform(b._2) }
-
-    logger.info(s"""----------------------------${"-" * name.size}---------\n""")
-    logger.info(f"Time: $timeMillis%.1f ms")
-
+    val result = Transform.runTransform(name, underlying.transform(b._2), logger)
     (b._1, result)
   }
-
-  /** Propagate annotations and update their names.
-    *
-    * @param inAnno input AnnotationSeq
-    * @param resAnno result AnnotationSeq
-    * @param renameOpt result RenameMap
-    * @return the updated annotations
-    */
-  private[firrtl] def propagateAnnotations(
-      inAnno: AnnotationSeq,
-      resAnno: AnnotationSeq,
-      renameOpt: Option[RenameMap]): AnnotationSeq = {
-    val newAnnotations = {
-      val inSet = mutable.LinkedHashSet() ++ inAnno
-      val resSet = mutable.LinkedHashSet() ++ resAnno
-      val deleted = (inSet -- resSet).map {
-        case DeletedAnnotation(xFormName, delAnno) => DeletedAnnotation(s"$xFormName+$name", delAnno)
-        case anno => DeletedAnnotation(name, anno)
-      }
-      val created = resSet -- inSet
-      val unchanged = resSet & inSet
-      (deleted ++ created ++ unchanged)
-    }
-
-    // For each annotation, rename all annotations.
-    val renames = renameOpt.getOrElse(RenameMap())
-    val remapped2original = mutable.LinkedHashMap[Annotation, mutable.LinkedHashSet[Annotation]]()
-    val keysOfNote = mutable.LinkedHashSet[Annotation]()
-    val finalAnnotations = newAnnotations.flatMap { anno =>
-      val remappedAnnos = anno.update(renames)
-      remappedAnnos.foreach { remapped =>
-        val set = remapped2original.getOrElseUpdate(remapped, mutable.LinkedHashSet.empty[Annotation])
-        set += anno
-        if(set.size > 1) keysOfNote += remapped
-      }
-      remappedAnnos
-    }.toSeq
-    keysOfNote.foreach { key =>
-      logger.debug(s"""The following original annotations are renamed to the same new annotation.""")
-      logger.debug(s"""Original Annotations:\n  ${remapped2original(key).mkString("\n  ")}""")
-      logger.debug(s"""New Annotation:\n  $key""")
-    }
-    finalAnnotations
-  }
 }
 
 object UpdateAnnotations {
diff --git a/src/test/scala/firrtlTests/annotationTests/JsonProtocolSpec.scala b/src/test/scala/firrtlTests/annotationTests/JsonProtocolSpec.scala
index 507c5291..d1bd1fb2 100644
--- a/src/test/scala/firrtlTests/annotationTests/JsonProtocolSpec.scala
+++ b/src/test/scala/firrtlTests/annotationTests/JsonProtocolSpec.scala
@@ -2,10 +2,12 @@
 
 package firrtlTests.annotationTests
 
-import firrtl.Parser
-import firrtl.annotations.{Annotation, JsonProtocol, NoTargetAnnotation}
+import firrtl._
+import firrtl.annotations.{JsonProtocol, NoTargetAnnotation}
 import firrtl.ir._
-import org.scalatest.{FlatSpec, Matchers, PropSpec}
+import firrtl.options.Dependency
+import _root_.logger.{Logger, LogLevel, LogLevelAnnotation}
+import org.scalatest.{FlatSpec, Matchers}
 
 case class AnAnnotation(
     info: Info,
@@ -18,6 +20,16 @@ case class AnAnnotation(
     groundType: GroundType
 ) extends NoTargetAnnotation
 
+class AnnoInjector extends Transform with DependencyAPIMigration {
+  override def optionalPrerequisiteOf = Dependency[ChirrtlEmitter] :: Nil
+  override def invalidates(a: Transform): Boolean = false
+  def execute(state: CircuitState): CircuitState = {
+    // Classes defined in method bodies can't be serialized by json4s
+    case class MyAnno(x: Int) extends NoTargetAnnotation
+    state.copy(annotations = MyAnno(3) +: state.annotations)
+  }
+}
+
 class JsonProtocolSpec extends FlatSpec with Matchers {
   "JsonProtocol" should "serialize and deserialize FIRRTL types" in {
 
@@ -40,4 +52,17 @@ class JsonProtocolSpec extends FlatSpec with Matchers {
     val outputAnnos = JsonProtocol.deserialize(annosString)
     inputAnnos should be (outputAnnos)
   }
+
+  "Annotation serialization during logging" should "not throw an exception" in {
+    val compiler = new firrtl.stage.transforms.Compiler(Seq(Dependency[AnnoInjector]))
+    val circuit = Parser.parse("""
+      |circuit test :
+      |  module test :
+      |    output out : UInt<1>
+      |    out <= UInt(0)
+      """.stripMargin)
+    Logger.makeScope(LogLevelAnnotation(LogLevel.Trace) :: Nil) {
+      compiler.execute(CircuitState(circuit, Nil))
+    }
+  }
 }