From 2e4f69fe066da2a7d9edab865cd5328320497b2a Mon Sep 17 00:00:00 2001
From: unlsycn <g.humphrey.sln@gmail.com>
Date: Mon, 23 Sep 2024 00:30:35 +0800
Subject: [PATCH] feat: implement SerializableModuleElaborator (#4409)

* feat: implement SerializableModuleElaborator

Signed-off-by: unlsycn <unlsycn@unlsycn.com>

* fix: return tuple of Readable from designImpl

Signed-off-by: unlsycn <unlsycn@unlsycn.com>

---------

Signed-off-by: unlsycn <unlsycn@unlsycn.com>
---
 .../util/SerializableModuleElaborator.scala   | 77 +++++++++++++++++++
 .../SerializableModuleElaboratorSpec.scala    | 43 +++++++++++
 2 files changed, 120 insertions(+)
 create mode 100644 src/main/scala/chisel3/experimental/util/SerializableModuleElaborator.scala
 create mode 100644 src/test/scala/chiselTests/experimental/SerializableModuleElaboratorSpec.scala

diff --git a/src/main/scala/chisel3/experimental/util/SerializableModuleElaborator.scala b/src/main/scala/chisel3/experimental/util/SerializableModuleElaborator.scala
new file mode 100644
index 00000000000..e60f50672ca
--- /dev/null
+++ b/src/main/scala/chisel3/experimental/util/SerializableModuleElaborator.scala
@@ -0,0 +1,77 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package chisel3.experimental.util
+
+import geny.Readable
+
+import chisel3.RawModule
+import chisel3.experimental.{SerializableModule, SerializableModuleGenerator, SerializableModuleParameter}
+
+import scala.reflect.runtime.universe
+import scala.reflect.runtime.universe.{runtimeMirror, typeOf}
+
+/** Mixin this trait to produce elaborators for [[SerializableModule]]
+  */
+trait SerializableModuleElaborator {
+
+  /**
+    * Implementation of a config API to serialize the [[SerializableModuleParameter]]
+    * @example
+    * {{{
+    *  def config(parameter: MySerializableModuleParameter) =
+    *    os.write.over(os.pwd / "config.json", configImpl(parameter))
+    * }}}
+    */
+  def configImpl[P <: SerializableModuleParameter: universe.TypeTag](
+    parameter: P
+  )(
+    implicit rwP: upickle.default.Writer[P]
+  ): Readable = upickle.default.write(parameter)
+
+  /**
+    * Implementation of a design API to elaborate [[SerializableModule]]
+    *
+    * @return A tuple of Readable, where the first is the firrtl and the second is the serializable annotations
+    * @example
+    * {{{
+    *  def design(parameter: os.Path) = {
+    *    val (firrtl, annos) = designImpl[MySerializableModule, MySerializableModuleParameter](os.read.stream(parameter))
+    *    os.write.over(os.pwd / "GCD.fir", firrtl)
+    *    os.write.over(os.pwd / "GCD.anno.json", annos)
+    *  }
+    * }}}
+    */
+  def designImpl[M <: SerializableModule[P]: universe.TypeTag, P <: SerializableModuleParameter: universe.TypeTag](
+    parameter: Readable
+  )(
+    implicit
+    rwP: upickle.default.Reader[P]
+  ): (Readable, Readable) = {
+    var fir: firrtl.ir.Circuit = null
+    val annos = Seq(
+      new chisel3.stage.phases.Elaborate,
+      new chisel3.stage.phases.Convert
+    ).foldLeft(
+      Seq(
+        chisel3.stage.ChiselGeneratorAnnotation(() =>
+          SerializableModuleGenerator(
+            runtimeMirror(getClass.getClassLoader)
+              .runtimeClass(typeOf[M].typeSymbol.asClass)
+              .asInstanceOf[Class[M]],
+            upickle.default.read[P](parameter)
+          ).module().asInstanceOf[RawModule]
+        )
+      ): firrtl.AnnotationSeq
+    ) { case (annos, stage) => stage.transform(annos) }
+      .flatMap {
+        case firrtl.stage.FirrtlCircuitAnnotation(circuit) =>
+          fir = circuit
+          None
+        case _: firrtl.options.Unserializable => None
+        case a => Some(a)
+      }
+    val firrtlStream: Readable = fir.serialize
+    val annoStream:   Readable = firrtl.annotations.JsonProtocol.serializeRecover(annos)
+    (firrtlStream, annoStream)
+  }
+}
diff --git a/src/test/scala/chiselTests/experimental/SerializableModuleElaboratorSpec.scala b/src/test/scala/chiselTests/experimental/SerializableModuleElaboratorSpec.scala
new file mode 100644
index 00000000000..39dffd090cc
--- /dev/null
+++ b/src/test/scala/chiselTests/experimental/SerializableModuleElaboratorSpec.scala
@@ -0,0 +1,43 @@
+package chiselTests
+package experimental
+
+import chisel3._
+import chiselTests.experimental.GCDSerializableModule
+import chisel3.experimental.util.SerializableModuleElaborator
+import geny.Readable
+import upickle.default.read
+
+class GCDSerializableModuleElaborator extends SerializableModuleElaborator {
+  val configPath = os.pwd / "config.json"
+  val firPath = os.pwd / "GCD.fir"
+  val annosPath = os.pwd / "GCD.anno.json"
+
+  def config(parameter: GCDSerializableModuleParameter) =
+    os.write.over(configPath, configImpl(parameter))
+
+  def design() = {
+    val (firrtl, annos) =
+      designImpl[GCDSerializableModule, GCDSerializableModuleParameter](os.read.stream(configPath))
+    os.write.over(firPath, firrtl)
+    os.write.over(annosPath, annos)
+  }
+}
+
+class SerializableModuleElaboratorSpec extends ChiselFlatSpec {
+  val elaborator = new GCDSerializableModuleElaborator
+  elaborator.config(GCDSerializableModuleParameter(16))
+  elaborator.design()
+
+  val firFile = os.read(elaborator.firPath)
+  val annosFile = os.read(elaborator.annosPath)
+
+  "SerializableModuleElaborator" should "elaborate firrtl" in {
+    firFile should include("module GCD")
+  }
+
+  "SerializableModuleElaborator" should "filter unserializable annotations" in {
+    (annosFile should not).include("UnserializeableAnnotation")
+    (annosFile should not).include("DesignAnnotation")
+    (annosFile should not).include("ChiselCircuitAnnotation")
+  }
+}