make PresetRegAnnotation public (#2254)

* make PresetRegAnnotation public

this annotation is useful outside the firrtl compiler:
- to implement a pass that creates registers which
  need to be initialized at the beginning of simulation
  (e.g., for formal verification)
- to support preset registers in treadle

* add PresetRegAnnotation test and deal with annotation correctly in RemoveReset pass

2 files changed, 41 insertions, 5 deletions

diff --git a/src/main/scala/firrtl/annotations/PresetAnnotations.scala b/src/main/scala/firrtl/annotations/PresetAnnotations.scala
index 32f24f13..449679cc 100644
--- a/src/main/scala/firrtl/annotations/PresetAnnotations.scala
+++ b/src/main/scala/firrtl/annotations/PresetAnnotations.scala
@@ -17,15 +17,31 @@ case class PresetAnnotation(target: ReferenceTarget)
 
 /**
   * Transform the targeted asynchronously-reset Reg into a bitstream preset Reg
-  * Used internally to annotate all registers associated to an AsyncReset tree
+  * Thus you can use this annotation in order to initialize a register
+  * at the beginning of simulation or through the FPGA bit-stream to its `init` value.
+  *
+  * The register must fulfil the following requirements:
+  * - the reset signal is `UInt(0)`
+  * - the `init` value is a Literal
   *
   * @param target ReferenceTarget to a Reg
   */
-private[firrtl] case class PresetRegAnnotation(
+case class PresetRegAnnotation(
   target: ReferenceTarget)
     extends SingleTargetAnnotation[ReferenceTarget]
-    with RegisterEmissionOption {
+    with RegisterEmissionOption
+    with firrtl.transforms.DontTouchAllTargets {
   def duplicate(n: ReferenceTarget) = this.copy(target = n)
   override def useInitAsPreset = true
   override def disableRandomization = true
 }
+
+object PresetRegAnnotation {
+
+  /** Extracts the names of every preset reg in the design by module. */
+  def collect(annotations: AnnotationSeq, main: String): Map[String, Set[String]] =
+    annotations.collect {
+      case a: PresetRegAnnotation if a.target.circuit == main =>
+        a.target.module -> a.target.ref
+    }.groupBy(_._1).map { case (k, v) => k -> v.map(_._2).toSet }
+}
diff --git a/src/main/scala/firrtl/transforms/RemoveReset.scala b/src/main/scala/firrtl/transforms/RemoveReset.scala
index 14880ab9..d7f59321 100644
--- a/src/main/scala/firrtl/transforms/RemoveReset.scala
+++ b/src/main/scala/firrtl/transforms/RemoveReset.scala
@@ -7,6 +7,7 @@ import firrtl.ir._
 import firrtl.Mappers._
 import firrtl.traversals.Foreachers._
 import firrtl.WrappedExpression.we
+import firrtl.annotations.PresetRegAnnotation
 import firrtl.options.Dependency
 
 import scala.collection.{immutable, mutable}
@@ -47,11 +48,26 @@ object RemoveReset extends Transform with DependencyAPIMigration {
     invalids.toSet
   }
 
-  private def onModule(m: DefModule): DefModule = {
+  private def onModule(m: DefModule, isPreset: String => Boolean): DefModule = {
     val resets = mutable.HashMap.empty[String, Reset]
     val invalids = computeInvalids(m)
     def onStmt(stmt: Statement): Statement = {
       stmt match {
+        case reg @ DefRegister(_, name, _, _, reset, init) if isPreset(name) =>
+          // registers that are preset annotated should already be in canonical form
+          if (reset != Utils.False()) {
+            throw new RuntimeException(
+              s"[${m.name}] register `$name` has a PresetRegAnnotation, but the reset is not UInt(0)!"
+            )
+          }
+          if (!Utils.isLiteral(init)) {
+            throw new RuntimeException(
+              s"[${m.name}] register `$name` has a PresetRegAnnotation, " +
+                s"but the init value is not a literal! ${init.serialize}"
+            )
+          }
+          // no change necessary
+          reg
         /* A register is initialized to an invalid expression */
         case reg @ DefRegister(_, _, _, _, _, init) if invalids.contains(we(init)) =>
           reg.copy(reset = Utils.zero, init = WRef(reg))
@@ -74,7 +90,11 @@ object RemoveReset extends Transform with DependencyAPIMigration {
   }
 
   def execute(state: CircuitState): CircuitState = {
-    val c = state.circuit.map(onModule)
+    // If registers are annotated with the [[PresetRegAnnotation]], they will take on their
+    // reset value when the circuit "starts" (i.e. at the beginning of simulation or when the FPGA
+    // bit-stream is initialized) and thus we need to special-case them.
+    val presetRegs = PresetRegAnnotation.collect(state.annotations, state.circuit.main)
+    val c = state.circuit.mapModule(m => onModule(m, presetRegs.getOrElse(m.name, _ => false)))
     state.copy(circuit = c)
   }
 }