Implement Lookupable for HasTarget (used by SRAM) (#4481)

Author: jackkoenig

core/src/main/scala/chisel3/experimental/hierarchy/core/Lookupable.scala

@@ -405,6 +405,46 @@ object Lookupable {
       }
     }
 
+  // TODO, this, cloneMemToContext, and cloneDataToContext should be unified
+  private def cloneHasTargetToContext(
+    hasTarget: HasTarget,
+    context:   BaseModule
+  )(
+    implicit sourceInfo: SourceInfo
+  ): HasTarget = {
+    hasTarget match {
+      case HasTarget.Impl(st: SramTarget) =>
+        st._parent match {
+          case None => hasTarget
+          case Some(parent) =>
+            val newParent = cloneModuleToContext(Proto(parent), context)
+            newParent match {
+              case Proto(p) if p == parent => hasTarget
+              case Clone(mod: BaseModule) =>
+                val existingMod = Builder.currentModule
+                Builder.currentModule = Some(mod)
+                val newChild = new SramTarget
+                Builder.currentModule = existingMod
+                newChild.setRef(st.getRef, true)
+                HasTarget(newChild)
+              case _ =>
+                throw new InternalErrorException(s"Match error: newParent=$newParent")
+            }
+        }
+    }
+  }
+
+  implicit def lookupHasTarget(implicit sourceInfo: SourceInfo): Simple[HasTarget] =
+    new Lookupable[HasTarget] {
+      type C = HasTarget
+      def definitionLookup[A](that: A => HasTarget, definition: Definition[A]): C = {
+        cloneHasTargetToContext(that(definition.proto), definition.getInnerDataContext.get)
+      }
+      def instanceLookup[A](that: A => HasTarget, instance: Instance[A]): C = {
+        cloneHasTargetToContext(that(instance.proto), instance.getInnerDataContext.get)
+      }
+    }
+
   import scala.language.higherKinds // Required to avoid warning for lookupIterable type parameter
   implicit def lookupIterable[B, F[_] <: Iterable[_]](
     implicit sourceInfo: SourceInfo,

core/src/main/scala/chisel3/package.scala

@@ -421,6 +421,8 @@ package object chisel3 {
 
   /** Exposes target information and suggestName functionality of a NamedComponent.
     */
+  // This is only currently used for SRAM to hide the underlying Memory but still let users annotate it.
+  // Rather than generalizing this, it's more likely that we'll just delete it (and the use in SRAM) in favor of Path Properties.
   sealed trait HasTarget {
     def toTarget:         ReferenceTarget
     def toAbsoluteTarget: ReferenceTarget
@@ -436,11 +438,7 @@ package object chisel3 {
   }
 
   object HasTarget {
-
-    /** This wrapping hides the actual object, ensuring users only have access
-      * to the target methods (instead of the type of the underlying object).
-      */
-    private[chisel3] def apply(t: NamedComponent): HasTarget = new HasTarget {
+    private[chisel3] case class Impl(t: SramTarget) extends HasTarget {
       def toTarget = t.toTarget
       def toAbsoluteTarget = t.toAbsoluteTarget
       def toRelativeTarget(root:            Option[BaseModule]) = t.toRelativeTarget(root)
@@ -449,5 +447,10 @@ package object chisel3 {
       def suggestName(seed: String): Unit = t.suggestName(seed)
     }
 
+    /** This wrapping hides the actual object, ensuring users only have access
+      * to the target methods (instead of the type of the underlying object).
+      */
+    private[chisel3] def apply(t: SramTarget): HasTarget = Impl(t)
+
   }
 }

src/test/scala/chiselTests/experimental/hierarchy/Annotations.scala

@@ -4,7 +4,7 @@ package chiselTests.experimental.hierarchy
 
 import _root_.firrtl.annotations._
 import chisel3.experimental.{annotate, BaseModule}
-import chisel3.{Data, MemBase}
+import chisel3.{Data, HasTarget, MemBase}
 import chisel3.experimental.hierarchy.{Definition, Hierarchy, Instance}
 
 // These annotations exist purely for testing purposes
@@ -24,8 +24,13 @@ private[hierarchy] object Annotations {
       extends chisel3.experimental.ChiselAnnotation {
     def toFirrtl = if (isAbsolute) MarkAnnotation(m.toAbsoluteTarget, tag) else MarkAnnotation(m.toTarget, tag)
   }
+  case class MarkChiselHasTargetAnnotation(d: HasTarget, tag: String, isAbsolute: Boolean)
+      extends chisel3.experimental.ChiselAnnotation {
+    def toFirrtl = if (isAbsolute) MarkAnnotation(d.toAbsoluteTarget, tag) else MarkAnnotation(d.toTarget, tag)
+  }
   def mark(d:                   Data, tag:         String): Unit = annotate(MarkChiselAnnotation(d, tag, false))
   def mark[T <: Data](d:        MemBase[T], tag:   String): Unit = annotate(MarkChiselMemAnnotation(d, tag, false))
+  def mark(d:                   HasTarget, tag:    String): Unit = annotate(MarkChiselHasTargetAnnotation(d, tag, false))
   def mark[B <: BaseModule](d:  Hierarchy[B], tag: String): Unit = annotate(MarkChiselHierarchyAnnotation(d, tag, true))
   def amark(d:                  Data, tag:         String): Unit = annotate(MarkChiselAnnotation(d, tag, true))
   def amark[B <: BaseModule](d: Hierarchy[B], tag: String): Unit = annotate(MarkChiselHierarchyAnnotation(d, tag, true))

src/test/scala/chiselTests/experimental/hierarchy/DefinitionSpec.scala

@@ -390,6 +390,16 @@ class DefinitionSpec extends ChiselFunSpec with Utils {
           "Cannot create a memory port in a different module (Top) than where the memory is (HasMems)."
       )
     }
+    it("(3.o): should work on HasTarget") {
+      class Top() extends Module {
+        val i = Definition(new HasHasTarget)
+        mark(i.x, "x")
+      }
+      val (_, annos) = getFirrtlAndAnnos(new Top)
+      annos.collect { case c: MarkAnnotation => c } should contain(
+        MarkAnnotation("~Top|HasHasTarget>sram_sram".rt, "x")
+      )
+    }
   }
   describe("(4): toDefinition") {
     it("(4.a): should work on modules") {

src/test/scala/chiselTests/experimental/hierarchy/Examples.scala

@@ -3,7 +3,7 @@
 package chiselTests.experimental.hierarchy
 
 import chisel3._
-import chisel3.util.Valid
+import chisel3.util.{SRAM, Valid}
 import chisel3.experimental.hierarchy._
 import chisel3.experimental.{attach, Analog, BaseModule}
 
@@ -223,6 +223,11 @@ object Examples {
     @public val xy = (x, y)
   }
   @instantiable
+  class HasHasTarget() extends Module {
+    val sram = SRAM(1024, UInt(8.W), 1, 1, 0)
+    @public val x: HasTarget = sram.underlying.get
+  }
+  @instantiable
   class HasVec() extends Module {
     @public val x = VecInit(1.U, 2.U, 3.U)
   }

src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala

@@ -464,6 +464,16 @@ class InstanceSpec extends ChiselFunSpec with Utils {
       val (chirrtl, _) = getFirrtlAndAnnos(new AddTwoNestedInstantiableDataWrapper(4))
       exactly(3, chirrtl.serialize.split('\n')) should include("connect i1.in, i0.out")
     }
+    it("(3.r): should work on HasTarget") {
+      class Top() extends Module {
+        val i = Instance(Definition(new HasHasTarget))
+        mark(i.x, "x")
+      }
+      val (_, annos) = getFirrtlAndAnnos(new Top)
+      annos.collect { case c: MarkAnnotation => c } should contain(
+        MarkAnnotation("~Top|Top/i:HasHasTarget>sram_sram".rt, "x")
+      )
+    }
   }
   describe("(4) toInstance") {
     it("(4.a): should work on modules") {