Refactor print control transform (#845)

src/main/scala/top/XiangShanStage.scala

@@ -17,74 +17,10 @@ package top
 
 import chisel3.stage.ChiselCli
 import firrtl.AnnotationSeq
-import firrtl.annotations.NoTargetAnnotation
-import firrtl.options.PhaseManager.PhaseDependency
 import firrtl.options.{Dependency, HasShellOptions, Shell, ShellOption}
 import firrtl.stage.{FirrtlCli, RunFirrtlTransformAnnotation}
-import freechips.rocketchip.stage.phases.GenerateArtefacts
 import freechips.rocketchip.transforms.naming.{OverrideDesiredNameAnnotation, RenameDesiredNames}
-import xstransforms.ShowPrintTransform
-import xstransforms.PrintModuleName
-
-case class DisablePrintfAnnotation(m: String) extends NoTargetAnnotation
-
-object DisablePrintfAnnotation extends HasShellOptions{
-
-  val options = Seq(
-    new ShellOption[String](
-      longOption = "disable-module-print",
-      toAnnotationSeq = s => Seq(DisablePrintfAnnotation(s)),
-      helpText =
-        "The verilog 'printf' in the <module> and it's submodules will be removed\n",
-      shortOption = Some("dm"),
-      helpValueName = Some("<module>")
-    )
-  )
-
-}
-
-case class EnablePrintfAnnotation(m: String) extends NoTargetAnnotation
-
-object EnablePrintfAnnotation extends HasShellOptions {
-  val options = Seq(
-    new ShellOption[String](
-      longOption = "enable-module-print",
-      toAnnotationSeq = s => Seq(EnablePrintfAnnotation(s)),
-      helpText =
-        "The verilog 'printf' except the <module> and it's submodules will be removed\n",
-      shortOption = Some("em"),
-      helpValueName = Some("<module>")
-    )
-  )
-
-}
-
-case class DisableAllPrintAnnotation() extends NoTargetAnnotation
-
-object DisableAllPrintAnnotation extends HasShellOptions {
-  val options = Seq(
-    new ShellOption[Unit](
-      longOption = "disable-all",
-      toAnnotationSeq = _ => Seq(DisableAllPrintAnnotation()),
-      helpText =
-        "All the verilog 'printf' will be removed\n",
-      shortOption = Some("dall")
-    )
-  )
-}
-
-case class RemoveAssertAnnotation() extends NoTargetAnnotation
-
-object RemoveAssertAnnotation extends HasShellOptions{
-  val options = Seq(
-    new ShellOption[Unit](
-      longOption = "remove-assert",
-      toAnnotationSeq = _ => Seq(RemoveAssertAnnotation()),
-      helpText = "All the 'assert' will be removed\n",
-      shortOption = None
-    )
-  )
-}
+import xstransforms._
 
 trait XiangShanCli { this: Shell =>
   parser.note("XiangShan Options")
@@ -110,7 +46,7 @@ object XiangShanStage {
     (new XiangShanStage).execute(
       args,
       annotations ++ Seq(
-        RunFirrtlTransformAnnotation(new ShowPrintTransform),
+        RunFirrtlTransformAnnotation(new PrintControl),
         RunFirrtlTransformAnnotation(new PrintModuleName),
         RunFirrtlTransformAnnotation(new RenameDesiredNames)
       )

src/main/scala/xstransforms/ShowPrintTransform.scala → src/main/scala/xstransforms/PrintControl.scala

@@ -16,37 +16,91 @@
 package xstransforms
 
 import firrtl._
+import firrtl.annotations.NoTargetAnnotation
 import firrtl.ir._
-import top._
+import firrtl.options.{HasShellOptions, ShellOption}
 
 import scala.collection.mutable
 
-class ShowPrintTransform extends Transform with DependencyAPIMigration {
+case class DisablePrintfAnnotation(m: String) extends NoTargetAnnotation
+object DisablePrintfAnnotation extends HasShellOptions{
+
+  val options = Seq(
+    new ShellOption[String](
+      longOption = "disable-module-print",
+      toAnnotationSeq = s => Seq(DisablePrintfAnnotation(s)),
+      helpText =
+        "The verilog 'printf' in the <module> and it's submodules will be removed\n",
+      shortOption = Some("dm"),
+      helpValueName = Some("<module>")
+    )
+  )
+
+}
+
+case class EnablePrintfAnnotation(m: String) extends NoTargetAnnotation
+object EnablePrintfAnnotation extends HasShellOptions {
+  val options = Seq(
+    new ShellOption[String](
+      longOption = "enable-module-print",
+      toAnnotationSeq = s => Seq(EnablePrintfAnnotation(s)),
+      helpText =
+        "The verilog 'printf' except the <module> and it's submodules will be removed\n",
+      shortOption = Some("em"),
+      helpValueName = Some("<module>")
+    )
+  )
+
+}
+
+case class DisableAllPrintAnnotation() extends NoTargetAnnotation
+object DisableAllPrintAnnotation extends HasShellOptions {
+  val options = Seq(
+    new ShellOption[Unit](
+      longOption = "disable-all",
+      toAnnotationSeq = _ => Seq(DisableAllPrintAnnotation()),
+      helpText =
+        "All the verilog 'printf' will be removed\n",
+      shortOption = Some("dall")
+    )
+  )
+}
+
+case class RemoveAssertAnnotation() extends NoTargetAnnotation
+object RemoveAssertAnnotation extends HasShellOptions{
+  val options = Seq(
+    new ShellOption[Unit](
+      longOption = "remove-assert",
+      toAnnotationSeq = _ => Seq(RemoveAssertAnnotation()),
+      helpText = "All the 'assert' will be removed\n",
+      shortOption = None
+    )
+  )
+}
+
+class PrintControl extends Transform with DependencyAPIMigration {
 
   override def optionalPrerequisiteOf = firrtl.stage.Forms.MinimalHighForm
   override def invalidates(a: Transform) = true
 
   override protected def execute(state: CircuitState): CircuitState = {
-    val c = state.circuit
 
-    val blackList = state.annotations.collect {
+    val disableList = state.annotations.collect {
       case DisablePrintfAnnotation(m) => m
     }
-    val whiteList = state.annotations.collect {
+    val enableList = state.annotations.collect {
       case EnablePrintfAnnotation(m) => m
     }
     val disableAll = state.annotations.collectFirst {
       case DisableAllPrintAnnotation() => true
     }.nonEmpty
-
     val removeAssert = state.annotations.collectFirst{
       case RemoveAssertAnnotation() => true
     }.nonEmpty
 
-    assert(
-      !(whiteList.nonEmpty && (disableAll || blackList.nonEmpty)),
-      "'white list' can't be used with 'disable all' or 'black list'!"
-    )
+    assert(!(enableList.nonEmpty && (disableAll || disableList.nonEmpty)))
+
+    val c = state.circuit
 
     val top = c.main
     val queue = new mutable.Queue[String]()
@@ -55,57 +109,39 @@ class ShowPrintTransform extends Transform with DependencyAPIMigration {
     queue += top
     ancestors(top) = mutable.LinkedHashSet.empty
 
-    def findSubModules(m: DefModule): Unit = {
-      def viewStmt(s: Statement): Statement = s match {
-        case DefInstance(_, name, module, _) =>
-          ancestors(module) = ancestors(m.name) + m.name
-          queue += module
-          s
-        case other =>
-          other.mapStmt(viewStmt)
-      }
-      m.foreachStmt(viewStmt)
-    }
-
     while (queue.nonEmpty) {
       val curr = queue.dequeue()
-      c.modules.find(m => m.name==curr).map(findSubModules)
+      c.modules.find(m => m.name==curr).foreach(m => {
+        def viewStmt(s: Statement): Statement = s match {
+          case DefInstance(_, _, module, _) =>
+            ancestors(module) = ancestors(curr) + m.name
+            queue += module
+            s
+          case other =>
+            other.mapStmt(viewStmt)
+        }
+        m.foreachStmt(viewStmt)
+      })
     }
 
-    def processModule(m: DefModule): DefModule = {
-      def disableModulePrint(mod: DefModule) = {
-        def disableStmtPrint(s: Statement): Statement = s match {
-          case _: Print =>
-            EmptyStmt
-          case other =>
-            other.mapStmt(disableStmtPrint)
+    def onModule(m: DefModule): DefModule = m match {
+      case _: ExtModule => m
+      case _: Module =>
+        def inRange(seq: Seq[String]): Boolean = {
+          seq.nonEmpty && (seq.contains(m.name) || seq.map(elm => {
+            ancestors(m.name).contains(elm)
+          }).reduce(_||_))
         }
-        mod.mapStmt(disableStmtPrint)
-      }
-      def removeModuleAssert(mod: DefModule)= {
-        def removeStmtAssert(s: Statement): Statement = s match {
-          case _: Stop =>
+        val enable = enableList.isEmpty || inRange(enableList)
+        val disable = disableAll || inRange(disableList) || !enable
+        def onStmt(s: Statement): Statement = s match {
+          case _: Print if disable =>
             EmptyStmt
-          case other =>
-            other.mapStmt(removeStmtAssert)
+          case _: Stop if removeAssert => EmptyStmt
+          case other => other.mapStmt(onStmt)
         }
-        mod.mapStmt(removeStmtAssert)
-      }
-
-      val isInBlackList = blackList.nonEmpty && (
-        blackList.contains(m.name) || blackList.map( b => ancestors(m.name).contains(b)).reduce(_||_)
-      )
-      val isInWhiteList = whiteList.isEmpty || (
-        whiteList.nonEmpty && (whiteList.contains(m.name) || whiteList.map( x => ancestors(m.name).contains(x)).reduce(_||_))
-      ) 
-      val tmpMod = if(disableAll || isInBlackList || !isInWhiteList){
-        disableModulePrint(m)
-      } else {
-        m
-      }
-      if(removeAssert) removeModuleAssert(tmpMod) else tmpMod
+        m.mapStmt(onStmt)
     }
-
-    state.copy(c.mapModule(processModule))
+    state.copy(circuit = c.mapModule(onModule))
   }
 }