Implement trait for Chisel compiler to name arbitrary non-Data types (#2610) (#2617)

Co-authored-by: Jack Koenig <koenig@sifive.com>
Co-authored-by: Megan Wachs <megan@sifive.com>
(cherry picked from commit 3ab34cddd8b87c22d5fc31020f10ddb2f1990d51)

Co-authored-by: Jared Barocsi <82000041+jared-barocsi@users.noreply.github.com>

4 files changed, 70 insertions, 28 deletions

diff --git a/core/src/main/scala/chisel3/experimental/package.scala b/core/src/main/scala/chisel3/experimental/package.scala
index 9b9c83f4..47b79099 100644
--- a/core/src/main/scala/chisel3/experimental/package.scala
+++ b/core/src/main/scala/chisel3/experimental/package.scala
@@ -162,6 +162,36 @@ package object experimental {
     */
   trait NoChiselNamePrefix
 
+  /** Generate prefixes from values of this type in the Chisel compiler plugin
+    *
+    * Users can mixin this trait to tell the Chisel compiler plugin to include the names of
+    * vals of this type when generating prefixes for naming `Data` and `Mem` instances.
+    * This is generally useful whenever creating a `class` that contains `Data`, `Mem`,
+    * or `Module` instances but does not itself extend `Data` or `Module`.
+    *
+    * @see See [[https://www.chisel-lang.org/chisel3/docs/explanations/naming.html the compiler plugin documentation]] for more information on this process.
+    *
+    * @example {{{
+    * import chisel3._
+    * import chisel3.experimental.AffectsChiselPrefix
+    *
+    * class MyModule extends Module {
+    *   // Note: This contains a Data but is not a named component itself
+    *   class NotAData extends AffectsChiselPrefix {
+    *     val value = Wire(Bool())
+    *   }
+    *
+    *   // Name with AffectsChiselPrefix:    "nonData_value"
+    *   // Name without AffectsChiselPrefix: "value"
+    *   val nonData = new NotAData
+    *
+    *   // Name with AffectsChiselPrefix:    "nonData2_value"
+    *   // Name without AffectsChiselPrefix: "value_1"
+    *   val nonData2 = new NotAData
+    * }
+    */
+  trait AffectsChiselPrefix
+
   object BundleLiterals {
     implicit class AddBundleLiteralConstructor[T <: Record](x: T) {
       def Lit(elems: (T => (Data, Data))*)(implicit sourceInfo: SourceInfo, compileOptions: CompileOptions): T = {
diff --git a/plugin/src/main/scala/chisel3/internal/plugin/ChiselComponent.scala b/plugin/src/main/scala/chisel3/internal/plugin/ChiselComponent.scala
index b3bbdbe4..dd9f24fb 100644
--- a/plugin/src/main/scala/chisel3/internal/plugin/ChiselComponent.scala
+++ b/plugin/src/main/scala/chisel3/internal/plugin/ChiselComponent.scala
@@ -82,9 +82,17 @@ class ChiselComponent(val global: Global, arguments: ChiselPluginArguments)
     // Checking for all chisel3.internal.NamedComponents, but since it is internal, we instead have
     // to match the public subtypes
     private val shouldMatchNamedComp: Type => Boolean =
-      shouldMatchGen(tq"chisel3.Data", tq"chisel3.MemBase[_]", tq"chisel3.VerificationStatement")
+      shouldMatchGen(
+        tq"chisel3.Data",
+        tq"chisel3.MemBase[_]",
+        tq"chisel3.VerificationStatement"
+      )
     private val shouldMatchModule:   Type => Boolean = shouldMatchGen(tq"chisel3.experimental.BaseModule")
     private val shouldMatchInstance: Type => Boolean = shouldMatchGen(tq"chisel3.experimental.hierarchy.Instance[_]")
+    private val shouldMatchChiselPrefixed: Type => Boolean =
+      shouldMatchGen(
+        tq"chisel3.experimental.AffectsChiselPrefix"
+      )
 
     // Given a type tree, infer the type and return it
     private def inferType(t: Tree): Type = localTyper.typed(t, nsc.Mode.TYPEmode).tpe
@@ -171,25 +179,38 @@ class ChiselComponent(val global: Global, arguments: ChiselPluginArguments)
       // Check if a subtree is a candidate
       case dd @ ValDef(mods, name, tpt, rhs) if okVal(dd) =>
         val tpe = inferType(tpt)
+        val isData = shouldMatchData(tpe)
+        val isNamedComp = isData || shouldMatchNamedComp(tpe)
+        val isPrefixed = isNamedComp || shouldMatchChiselPrefixed(tpe)
+
         // If a Data and in a Bundle, just get the name but not a prefix
-        if (shouldMatchData(tpe) && inBundle(dd)) {
+        if (isData && inBundle(dd)) {
           val str = stringFromTermName(name)
           val newRHS = transform(rhs) // chisel3.internal.plugin.autoNameRecursively
           val named = q"chisel3.internal.plugin.autoNameRecursively($str)($newRHS)"
           treeCopy.ValDef(dd, mods, name, tpt, localTyper.typed(named))
         }
         // If a Data or a Memory, get the name and a prefix
-        else if (shouldMatchNamedComp(tpe)) {
+        else if (isData || isPrefixed) {
           val str = stringFromTermName(name)
           // Starting with '_' signifies a temporary, we ignore it for prefixing because we don't
           // want double "__" in names when the user is just specifying a temporary
           val prefix = if (str.head == '_') str.tail else str
           val newRHS = transform(rhs)
           val prefixed = q"chisel3.experimental.prefix.apply[$tpt](name=$prefix)(f=$newRHS)"
-          val named = q"chisel3.internal.plugin.autoNameRecursively($str)($prefixed)"
+
+          val named =
+            if (isNamedComp) {
+              // Only name named components (not things that are merely prefixed)
+              q"chisel3.internal.plugin.autoNameRecursively($str)($prefixed)"
+            } else {
+              prefixed
+            }
+
           treeCopy.ValDef(dd, mods, name, tpt, localTyper.typed(named))
-          // If an instance, just get a name but no prefix
-        } else if (shouldMatchModule(tpe)) {
+        }
+        // If an instance, just get a name but no prefix
+        else if (shouldMatchModule(tpe)) {
           val str = stringFromTermName(name)
           val newRHS = transform(rhs)
           val named = q"chisel3.internal.plugin.autoNameRecursively($str)($newRHS)"
diff --git a/src/main/scala/chisel3/util/Counter.scala b/src/main/scala/chisel3/util/Counter.scala
index ef1eff9f..0d1b8db1 100644
--- a/src/main/scala/chisel3/util/Counter.scala
+++ b/src/main/scala/chisel3/util/Counter.scala
@@ -3,6 +3,7 @@
 package chisel3.util
 
 import chisel3._
+import chisel3.experimental.AffectsChiselPrefix
 import chisel3.internal.naming.chiselName // can't use chisel3_ version because of compile order
 
 /** Used to generate an inline (logic directly in the containing Module, no internal Module is created)
@@ -27,8 +28,7 @@ import chisel3.internal.naming.chiselName // can't use chisel3_ version because
   *   }
   * }}}
   */
-@chiselName
-class Counter private (r: Range, oldN: Option[Int] = None) {
+class Counter private (r: Range, oldN: Option[Int] = None) extends AffectsChiselPrefix {
   require(r.length > 0, s"Counter range cannot be empty, got: $r")
   require(r.start >= 0 && r.end >= 0, s"Counter range must be positive, got: $r")
 
diff --git a/src/test/scala/chiselTests/NamingAnnotationTest.scala b/src/test/scala/chiselTests/NamingAnnotationTest.scala
index ded321cd..a3f39c51 100644
--- a/src/test/scala/chiselTests/NamingAnnotationTest.scala
+++ b/src/test/scala/chiselTests/NamingAnnotationTest.scala
@@ -4,12 +4,13 @@ package chiselTests
 
 import chisel3._
 import chisel3.experimental.chiselName
+import chisel3.experimental.AffectsChiselPrefix
 import chisel3.internal.InstanceId
 import chisel3.stage.ChiselStage
 
 import scala.collection.mutable.ListBuffer
 
-trait NamedModuleTester extends Module {
+trait NamedModuleTester extends Module with AffectsChiselPrefix {
   val expectedNameMap = ListBuffer[(InstanceId, String)]()
   val expectedModuleNameMap = ListBuffer[(Module, String)]()
 
@@ -48,25 +49,20 @@ trait NamedModuleTester extends Module {
     failures.toList
   }
 }
-@chiselName
-class OuterNamedNonModule {
+class OuterNamedNonModule extends AffectsChiselPrefix {
   val value = Wire(Bool())
 }
 
-@chiselName
-class NonModule {
+class NonModule extends AffectsChiselPrefix {
   val value = Wire(Bool())
-  @chiselName
-  class InnerNamedNonModule {
+  class InnerNamedNonModule extends AffectsChiselPrefix {
     val value = Wire(Bool())
   }
   val inner = new InnerNamedNonModule
   val outer = new OuterNamedNonModule
 }
 
-@chiselName
 class NamedModule extends NamedModuleTester {
-  @chiselName
   def FunctionMockupInner(): UInt = {
     val my2A = 1.U
     val my2B = expectName(my2A +& 2.U, "test_myNested_my2B")
@@ -74,7 +70,6 @@ class NamedModule extends NamedModuleTester {
     my2C
   }
 
-  @chiselName
   def FunctionMockup(): UInt = {
     val myNested = expectName(FunctionMockupInner(), "test_myNested")
     val myA = expectName(1.U + myNested, "test_myA")
@@ -123,11 +118,9 @@ class NamedModule extends NamedModuleTester {
   NoReturnFunction()
 }
 
-@chiselName
 class NameCollisionModule extends NamedModuleTester {
-  @chiselName
-  def repeatedCalls(id: Int): UInt = {
-    val test = expectName(1.U + 3.U, s"test_$id") // should disambiguate by invocation order
+  def repeatedCalls(name: String): UInt = {
+    val test = expectName(1.U + 3.U, s"${name}_test") // should disambiguate by invocation order
     test + 2.U
   }
 
@@ -135,8 +128,8 @@ class NameCollisionModule extends NamedModuleTester {
   def innerNamedFunction() {
     // ... but not this inner function
     def innerUnnamedFunction() {
-      val a = repeatedCalls(1)
-      val b = repeatedCalls(2)
+      val a = repeatedCalls("a")
+      val b = repeatedCalls("b")
     }
 
     innerUnnamedFunction()
@@ -212,19 +205,17 @@ class NoChiselNamePrefixTester extends NamedModuleTester {
     val a = expectName(1.U +& 2.U, "a")
   }
   val inst = new NoChiselNamePrefixClass
-  @chiselName
   class NormalClass {
     val b = 1.U +& 2.U
   }
-  val foo = new NormalClass
+  val foo = new NormalClass with AffectsChiselPrefix
   expectName(foo.b, "foo_b")
   val bar = new NormalClass with chisel3.experimental.NoChiselNamePrefix
   expectName(bar.b, "b")
 
   // Check that we're not matching by name but actual type
   trait NoChiselNamePrefix
-  @chiselName
-  class FakeNoChiselNamePrefix extends NoChiselNamePrefix {
+  class FakeNoChiselNamePrefix extends NoChiselNamePrefix with AffectsChiselPrefix {
     val c = 1.U +& 2.U
   }
   val fizz = new FakeNoChiselNamePrefix