Implement CustomRadixTransform for wave viewers (#2434)

1. Add CustomRadix{Def,Apply}Annotation to define and apply custom radix.
2. Add CustomRadixConfigFileAnnotation to output a JSON config file so
   users can generate scripts on their own.

Reviewed-by: Jiuyang Liu <liu@jiuyang.me>
Co-authored-by: sinofp <sinofp@tuta.io>

3 files changed, 247 insertions, 2 deletions

diff --git a/src/main/scala/firrtl/stage/FirrtlCli.scala b/src/main/scala/firrtl/stage/FirrtlCli.scala
index a9c08627..d416e25f 100644
--- a/src/main/scala/firrtl/stage/FirrtlCli.scala
+++ b/src/main/scala/firrtl/stage/FirrtlCli.scala
@@ -4,7 +4,7 @@ package firrtl.stage
 
 import firrtl.options.Shell
 import firrtl.passes.CommonSubexpressionElimination
-import firrtl.transforms.NoCircuitDedupAnnotation
+import firrtl.transforms.{CustomRadixTransform, NoCircuitDedupAnnotation}
 
 /** [[firrtl.options.Shell Shell]] mixin that provides command line options for FIRRTL. This does not include any
   * [[firrtl.options.RegisteredLibrary RegisteredLibrary]] or [[firrtl.options.RegisteredTransform RegisteredTransform]]
@@ -28,7 +28,8 @@ trait FirrtlCli { this: Shell =>
     OptimizeForFPGA,
     CurrentFirrtlStateAnnotation,
     CommonSubexpressionElimination,
-    AllowUnrecognizedAnnotations
+    AllowUnrecognizedAnnotations,
+    CustomRadixTransform
   )
     .map(_.addOptions(parser))
 
diff --git a/src/main/scala/firrtl/transforms/CustomRadixTransform.scala b/src/main/scala/firrtl/transforms/CustomRadixTransform.scala
new file mode 100644
index 00000000..42724be8
--- /dev/null
+++ b/src/main/scala/firrtl/transforms/CustomRadixTransform.scala
@@ -0,0 +1,115 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package firrtl.transforms
+
+import firrtl.annotations.TargetToken.Instance
+import firrtl.annotations.{Annotation, NoTargetAnnotation, ReferenceTarget, SingleTargetAnnotation}
+import firrtl.options.{CustomFileEmission, Dependency, HasShellOptions, ShellOption}
+import firrtl.stage.TransformManager.TransformDependency
+import firrtl.stage.{Forms, RunFirrtlTransformAnnotation}
+import firrtl.{AnnotationSeq, CircuitState, DependencyAPIMigration, Transform}
+
+/** Contains a static map from signal value(BigInt) to signal name(String)
+  * This is useful for enumeration(finite state machine, bus transaction name, etc)
+  *
+  * @param name identifier for this alias
+  * @param filters a sequence of translation filter
+  * @param width width of this alias
+  */
+case class CustomRadixDefAnnotation(name: String, filters: Seq[(BigInt, String)], width: Int) extends NoTargetAnnotation
+
+/** A annotation making a ReferenceTarget to be a specific [[CustomRadixDefAnnotation]].
+  *
+  * @param target the ReferenceTarget which the alias applied to
+  * @param name the identifier for the alias
+  */
+case class CustomRadixApplyAnnotation(target: ReferenceTarget, name: String)
+    extends SingleTargetAnnotation[ReferenceTarget] {
+  override def duplicate(n: ReferenceTarget): Annotation = this.copy(n)
+}
+
+/** Dumps a JSON config file for custom radix. Users can generate script using the emitted file.
+  *
+  * @param signals which alias contains which signals, the signals should be converted from ReferenceTarget to String
+  * @param filters sequence of [[CustomRadixDefAnnotation]], the name should match [[signals]].map(_._1)
+  */
+case class CustomRadixConfigFileAnnotation(
+  signals: Seq[(String, Seq[String])],
+  filters: Seq[CustomRadixDefAnnotation])
+    extends NoTargetAnnotation
+    with CustomFileEmission {
+  def waveViewer = "config"
+  def baseFileName(annotations: AnnotationSeq): String = "custom_radix"
+  def suffix: Option[String] = Some(".json")
+
+  def getBytes: Iterable[Byte] = {
+    import org.json4s.JsonDSL.WithBigDecimal._
+    import org.json4s.native.JsonMethods._
+    val aliasMap = signals.toMap
+    pretty(
+      render(
+        filters.map { a =>
+          val values = a.filters.map {
+            case (int, str) =>
+              ("digit" -> int) ~
+                ("alias" -> str)
+          }
+          a.name -> (
+            ("width" -> a.width) ~
+              ("values" -> values) ~
+              ("signals" -> aliasMap(a.name))
+          )
+        }
+      )
+    )
+  }.getBytes
+}
+
+/** A Transform that generate scripts or config file for Custom Radix */
+object CustomRadixTransform extends Transform with DependencyAPIMigration with HasShellOptions {
+
+  val options = Seq(
+    new ShellOption[String](
+      longOption = "wave-viewer-script",
+      toAnnotationSeq = str => {
+        RunFirrtlTransformAnnotation(Dependency(CustomRadixTransform)) +: str.toLowerCase
+          .split(',')
+          .map {
+            case "json" | "" => CustomRadixConfigFileAnnotation(Seq.empty, Seq.empty)
+          }
+          .toSeq
+      },
+      helpText = "<json>, you can combine them like 'json', pass empty string will generate json",
+      shortOption = None
+    )
+  )
+
+  // in case of any rename during transforms.
+  override def optionalPrerequisites: Seq[TransformDependency] = Forms.BackendEmitters
+  override def invalidates(a: Transform) = false
+
+  protected def execute(state: CircuitState): CircuitState = {
+    val annos = state.annotations
+    def toVerilogName(target: ReferenceTarget) =
+      (target.circuit +: target.tokens.collect { case Instance(i) => i } :+ target.ref).mkString(".")
+    // todo if scalaVersion >= 2.13: use groupMap
+    val filters = annos.collect { case a: CustomRadixDefAnnotation => a }
+
+    val signals = annos.collect {
+      case CustomRadixApplyAnnotation(target, name) =>
+        assert(filters.exists(_.name == name), s"$name not found in CustomRadixDefAnnotation.")
+        name -> target
+    }
+      .groupBy(_._1)
+      .mapValues(_.map(t => toVerilogName(t._2)))
+      .toSeq
+      .sortBy(_._1)
+
+    assert(filters.groupBy(_.name).forall(_._2.length == 1), "name collision in CustomRadixDefAnnotation detected.")
+
+    state.copy(annotations = annos.map {
+      case _: CustomRadixConfigFileAnnotation => CustomRadixConfigFileAnnotation(signals, filters)
+      case a => a
+    })
+  }
+}
diff --git a/src/test/scala/firrtlTests/transforms/CustomRadixTransformSpec.scala b/src/test/scala/firrtlTests/transforms/CustomRadixTransformSpec.scala
new file mode 100644
index 00000000..20d685dc
--- /dev/null
+++ b/src/test/scala/firrtlTests/transforms/CustomRadixTransformSpec.scala
@@ -0,0 +1,129 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package firrtlTests.transforms
+
+import firrtl.annotations.ReferenceTarget
+import firrtl.annotations.TargetToken.{Instance, OfModule}
+import firrtl.testutils.FirrtlFlatSpec
+import firrtl.stage.{FirrtlCircuitAnnotation, FirrtlStage}
+import firrtl.transforms.{CustomRadixApplyAnnotation, CustomRadixDefAnnotation}
+import firrtl.util.BackendCompilationUtilities
+
+class CustomRadixTransformSpec extends FirrtlFlatSpec {
+  behavior.of("CustomRadix")
+
+  val testDir = os.Path(BackendCompilationUtilities.createTestDirectory("CustomRadix").getAbsolutePath)
+  val fir =
+    """circuit M2 :
+      |  module PT :
+      |    input clock : Clock
+      |    input reset : Reset
+      |    output io : { flip i : UInt<7>, o : UInt<7>}
+      |
+      |    io.o <= io.i
+      |
+      |  module PT_1 :
+      |    input clock : Clock
+      |    input reset : Reset
+      |    output io : { flip i : UInt<7>, o : UInt<7>}
+      |
+      |    io.o <= io.i
+      |
+      |  module M1 :
+      |    input clock : Clock
+      |    input reset : Reset
+      |    output out : UInt<7>
+      |
+      |    reg cnt : UInt<5>, clock with :
+      |      reset => (reset, UInt<5>("h0"))
+      |    node _cnt_T = add(cnt, UInt<1>("h1"))
+      |    node _cnt_T_1 = tail(_cnt_T, 1)
+      |    cnt <= _cnt_T_1
+      |    inst pt of PT
+      |    pt.clock <= clock
+      |    pt.reset <= reset
+      |    inst pt2 of PT_1
+      |    pt2.clock <= clock
+      |    pt2.reset <= reset
+      |    pt2.io.i <= pt.io.o
+      |    pt2.io.o is invalid
+      |    pt.io.i <= UInt<1>("h0")
+      |    node _T = eq(cnt, UInt<1>("h1"))
+      |    when _T :
+      |      pt.io.i <= UInt<1>("h1")
+      |    else :
+      |      node _T_1 = eq(cnt, UInt<2>("h2"))
+      |      when _T_1 :
+      |        pt.io.i <= UInt<2>("h2")
+      |      else :
+      |        node _T_2 = eq(cnt, UInt<2>("h3"))
+      |        when _T_2 :
+      |          pt.io.i <= UInt<7>("h64")
+      |        else :
+      |          node _T_3 = eq(cnt, UInt<3>("h4"))
+      |          when _T_3 :
+      |            pt.io.i <= UInt<7>("h65")
+      |    out <= pt.io.o
+      |
+      |  module PT_2 :
+      |    input clock : Clock
+      |    input reset : Reset
+      |    output io : { flip i : UInt<7>, o : UInt<7>}
+      |
+      |    io.o <= io.i
+      |
+      |  module M2 :
+      |    input clock : Clock
+      |    input reset : UInt<1>
+      |    output out : UInt<7>
+      |
+      |    inst m1 of M1
+      |    m1.clock <= clock
+      |    m1.reset <= reset
+      |    inst pt3 of PT_2
+      |    pt3.clock <= clock
+      |    pt3.reset <= reset
+      |    pt3.io.i <= m1.out
+      |    out <= pt3.io.o
+      |""".stripMargin
+
+  val annotations = Seq(
+    FirrtlCircuitAnnotation(firrtl.Parser.parse(fir)),
+    CustomRadixDefAnnotation("EnumExample", Seq(0, 1, 2, 100, 101).map(x => BigInt(x) -> s"e$x"), 7)
+  ) ++ Seq(
+    ("M1", Seq(), "in"),
+    ("M2", Seq(Instance("pt3") -> OfModule("PT")), "io_o"),
+    ("M2", Seq(Instance("m1") -> OfModule("M1"), Instance("pt2") -> OfModule("PT")), "io_i")
+  ).map {
+    case (module, path, ref) =>
+      CustomRadixApplyAnnotation(ReferenceTarget("M2", module, path, ref, Seq()), "EnumExample")
+  }
+
+  it should "generate a JSON config file" in {
+    (new FirrtlStage).execute(Array("--wave-viewer-script", "", "--target-dir", testDir.toString), annotations)
+    val expected =
+      """[{
+        |  "EnumExample":{
+        |    "width":7,
+        |    "values":[{
+        |      "digit":0,
+        |      "alias":"e0"
+        |    },{
+        |      "digit":1,
+        |      "alias":"e1"
+        |    },{
+        |      "digit":2,
+        |      "alias":"e2"
+        |    },{
+        |      "digit":100,
+        |      "alias":"e100"
+        |    },{
+        |      "digit":101,
+        |      "alias":"e101"
+        |    }],
+        |    "signals":["M2.m1.pt.io_i","M2.m1.pt.io_o","M2.in"]
+        |  }
+        |}]""".stripMargin
+    assert(expected == os.read(testDir / "custom_radix.json"))
+  }
+}