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

3 files changed, 122 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)
   }
 }
diff --git a/src/test/scala/firrtlTests/PresetRegAnnotationSpec.scala b/src/test/scala/firrtlTests/PresetRegAnnotationSpec.scala
new file mode 100644
index 00000000..5cd0f764
--- /dev/null
+++ b/src/test/scala/firrtlTests/PresetRegAnnotationSpec.scala
@@ -0,0 +1,81 @@
+package firrtlTests
+
+import firrtl._
+import firrtl.annotations.{CircuitTarget, PresetRegAnnotation}
+import firrtl.options.Dependency
+import firrtl.testutils.LeanTransformSpec
+import firrtl.transforms.PropagatePresetAnnotations
+import logger.{LogLevel, LogLevelAnnotation, Logger}
+
+import scala.collection.mutable
+
+/** Tests the use of the [[firrtl.annotations.PresetRegAnnotation]]
+  * from a pass that needs to create a register with an initial value.
+  */
+class PresetRegAnnotationSpec
+    extends LeanTransformSpec(Seq(Dependency(MakePresetRegs), Dependency[SystemVerilogEmitter])) {
+  behavior.of("PresetRegAnnotation")
+
+  val src =
+    """circuit test:
+      |  module test:
+      |    input clock : Clock
+      |    output out : UInt<8>
+      |
+      |    reg r : UInt<8>, clock with :
+      |      reset => (UInt(0), UInt<8>("h7b"))
+      |    out <= r
+      |""".stripMargin
+
+  val ll = LogLevel.Error
+
+  it should "allow passes to mark registers to be initialized" in {
+    val result = Logger.makeScope(Seq(LogLevelAnnotation(ll))) { compile(src) }
+    val verilog = result.getEmittedCircuit.value
+    val lines = verilog.split('\n').map(_.trim).toSet
+
+    // the register should be assigned its initial value
+    assert(lines.contains("reg [7:0] r = 8'h7b;"))
+    assert(lines.contains("assign out = r;"))
+
+    // no asynchronous reset should be emitted
+    assert(!lines.contains("if (reset) begin"))
+    assert(!lines.contains("always @(posedge clock or posedge reset) begin"))
+  }
+}
+
+/** Adds preset reg annotations for all registers that have:
+  * 1. reset = false
+  * 2. init = a literal
+  */
+private object MakePresetRegs extends Transform with DependencyAPIMigration {
+  // run on lowered firrtl
+  override def prerequisites = Seq(Dependency(firrtl.passes.ExpandWhens), Dependency(firrtl.passes.LowerTypes))
+  override def invalidates(a: Transform) = false
+  // since we generate PresetRegAnnotations, we need to run after preset propagation
+  override def optionalPrerequisites = Seq(Dependency[PropagatePresetAnnotations])
+  // we want to run before the actual Verilog is emitted
+  // we want to look at the reset value, which may be removed by the RemoveReset transform.
+  override def optionalPrerequisiteOf = Seq(Dependency[SystemVerilogEmitter], Dependency(firrtl.transforms.RemoveReset))
+
+  override def execute(state: CircuitState): CircuitState = {
+    val c = CircuitTarget(state.circuit.main)
+    val newAnnos = state.circuit.modules.flatMap(onModule(c, _))
+    state.copy(annotations = newAnnos ++: state.annotations)
+  }
+
+  private def onModule(c: CircuitTarget, m: ir.DefModule): List[PresetRegAnnotation] = m match {
+    case mod: ir.Module =>
+      val regs = mutable.ListBuffer[String]()
+      mod.foreachStmt(onStmt(_, regs))
+      val m = c.module(mod.name)
+      regs.map(r => PresetRegAnnotation(m.ref(r))).toList
+    case _ => List()
+  }
+
+  private def onStmt(s: ir.Statement, regs: mutable.ListBuffer[String]): Unit = s match {
+    case ir.DefRegister(_, name, _, _, reset, init) if reset == Utils.False() && Utils.isLiteral(init) =>
+      regs.append(name)
+    case other => other.foreachStmt(onStmt(_, regs))
+  }
+}