Hierarchy API: make Mems lookupable

core/src/main/scala/chisel3/Mem.scala

@@ -246,6 +246,11 @@ sealed abstract class MemBase[T <: Data](val t: T, val length: BigInt)
   )(
     implicit compileOptions: CompileOptions
   ): T = {
+    if (Builder.currentModule != _parent) {
+      throwException(
+        s"Cannot create a memory port in a different module (${Builder.currentModule.get.name}) than where the memory is (${_parent.get.name})."
+      )
+    }
     requireIsHardware(idx, "memory port index")
     val i = Vec.truncateIndex(idx, length)(sourceInfo, compileOptions)
 
@@ -267,7 +272,7 @@ sealed abstract class MemBase[T <: Data](val t: T, val length: BigInt)
   * @note when multiple conflicting writes are performed on a Mem element, the
   * result is undefined (unlike Vec, where the last assignment wins)
   */
-sealed class Mem[T <: Data] private (t: T, length: BigInt) extends MemBase(t, length)
+sealed class Mem[T <: Data] private[chisel3] (t: T, length: BigInt) extends MemBase(t, length)
 
 object SyncReadMem {
 
@@ -345,7 +350,7 @@ object SyncReadMem {
   * @note when multiple conflicting writes are performed on a Mem element, the
   * result is undefined (unlike Vec, where the last assignment wins)
   */
-sealed class SyncReadMem[T <: Data] private (t: T, n: BigInt, val readUnderWrite: SyncReadMem.ReadUnderWrite)
+sealed class SyncReadMem[T <: Data] private[chisel3] (t: T, n: BigInt, val readUnderWrite: SyncReadMem.ReadUnderWrite)
     extends MemBase[T](t, n) {
 
   override def read(x: UInt): T = macro SourceInfoTransform.xArg

core/src/main/scala/chisel3/Module.scala

@@ -101,8 +101,8 @@ object Module extends SourceInfoDoc {
     compileOptions:      CompileOptions
   ): T = {
     val parent = Builder.currentModule
-
     val module: T = bc // bc is actually evaluated here
+    if (!parent.isEmpty) { Builder.currentModule = parent }
 
     module
   }
@@ -229,6 +229,8 @@ package internal {
     // Private internal class to serve as a _parent for Data in cloned ports
     private[chisel3] class ModuleClone[T <: BaseModule](val getProto: T) extends PseudoModule with IsClone[T] {
       override def toString = s"ModuleClone(${getProto})"
+      // Do not call default addId function, which may modify a module that is already "closed"
+      override def addId(d: HasId): Unit = ()
       def getPorts = _portsRecord
       // ClonePorts that hold the bound ports for this module
       // Used for setting the refs of both this module and the Record
@@ -307,6 +309,8 @@ package internal {
       override def toString = s"DefinitionClone(${getProto})"
       // No addition components are generated
       private[chisel3] def generateComponent(): Option[Component] = None
+      // Do not call default addId function, which may modify a module that is already "closed"
+      override def addId(d: HasId): Unit = ()
       // Necessary for toTarget to work
       private[chisel3] def initializeInParent(parentCompileOptions: CompileOptions): Unit = ()
       // Module name is the same as proto's module name

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

@@ -19,7 +19,7 @@ import chisel3.internal.{throwException, AggregateViewBinding, Builder, ChildBin
   */
 @implicitNotFound(
   "@public is only legal within a class or trait marked @instantiable, and only on vals of type" +
-    " Data, BaseModule, IsInstantiable, IsLookupable, or Instance[_], or in an Iterable, Option, Either, or Tuple2"
+    " Data, BaseModule, MemBase, IsInstantiable, IsLookupable, or Instance[_], or in an Iterable, Option, Either, or Tuple2"
 )
 trait Lookupable[-B] {
   type C // Return type of the lookup
@@ -348,6 +348,45 @@ object Lookupable {
       }
     }
 
+  private[chisel3] def cloneMemToContext[T <: MemBase[_]](
+    mem:     T,
+    context: BaseModule
+  )(
+    implicit sourceInfo: SourceInfo,
+    compileOptions:      CompileOptions
+  ): T = {
+    mem._parent match {
+      case None => mem
+      case Some(parent) =>
+        val newParent = cloneModuleToContext(Proto(parent), context)
+        newParent match {
+          case Proto(p) if p == parent => mem
+          case Clone(mod: BaseModule) =>
+            val existingMod = Builder.currentModule
+            Builder.currentModule = Some(mod)
+            val newChild: T = mem match {
+              case m: Mem[_] => new Mem(m.t.asInstanceOf[Data].cloneTypeFull, m.length).asInstanceOf[T]
+              case m: SyncReadMem[_] =>
+                new SyncReadMem(m.t.asInstanceOf[Data].cloneTypeFull, m.length, m.readUnderWrite).asInstanceOf[T]
+            }
+            Builder.currentModule = existingMod
+            newChild.setRef(mem.getRef, true)
+            newChild
+        }
+    }
+  }
+
+  implicit def lookupMem[B <: MemBase[_]](implicit sourceInfo: SourceInfo, compileOptions: CompileOptions) =
+    new Lookupable[B] {
+      type C = B
+      def definitionLookup[A](that: A => B, definition: Definition[A]): C = {
+        cloneMemToContext(that(definition.proto), definition.getInnerDataContext.get)
+      }
+      def instanceLookup[A](that: A => B, instance: Instance[A]): C = {
+        cloneMemToContext(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/internal/Builder.scala

@@ -258,6 +258,7 @@ private[chisel3] trait HasId extends InstanceId {
       (p._component, this) match {
         case (Some(c), _) => refName(c)
         case (None, d: Data) if d.topBindingOpt == Some(CrossModuleBinding) => _ref.get.localName
+        case (None, _: MemBase[Data]) => _ref.get.localName
         case (None, _) =>
           throwException(s"signalName/pathName should be called after circuit elaboration: $this, ${_parent}")
       }

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

@@ -4,10 +4,11 @@ package chiselTests.experimental.hierarchy
 
 import _root_.firrtl.annotations._
 import chisel3.experimental.{annotate, BaseModule}
-import chisel3.Data
+import chisel3.{Data, MemBase}
 import chisel3.experimental.hierarchy.{Definition, Hierarchy, Instance}
 
-object Annotations {
+// These annotations exist purely for testing purposes
+private[hierarchy] object Annotations {
   case class MarkAnnotation(target: IsMember, tag: String) extends SingleTargetAnnotation[IsMember] {
     def duplicate(n: IsMember): Annotation = this.copy(target = n)
   }
@@ -19,7 +20,12 @@ object Annotations {
       extends chisel3.experimental.ChiselAnnotation {
     def toFirrtl = if (isAbsolute) MarkAnnotation(d.toAbsoluteTarget, tag) else MarkAnnotation(d.toTarget, tag)
   }
+  case class MarkChiselMemAnnotation[T <: Data](m: MemBase[T], tag: String, isAbsolute: Boolean)
+      extends chisel3.experimental.ChiselAnnotation {
+    def toFirrtl = if (isAbsolute) MarkAnnotation(m.toAbsoluteTarget, tag) else MarkAnnotation(m.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[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

@@ -339,6 +339,29 @@ class DefinitionSpec extends ChiselFunSpec with Utils {
       annos should contain(MarkAnnotation("~Top|HasTuple2>x".rt, "x"))
       annos should contain(MarkAnnotation("~Top|HasTuple2>y".rt, "y"))
     }
+    it("3.13: should work on Mems/SyncReadMems") {
+      class Top() extends Module {
+        val i = Definition(new HasMems())
+        mark(i.mem, "Mem")
+        mark(i.syncReadMem, "SyncReadMem")
+      }
+      val (_, annos) = getFirrtlAndAnnos(new Top)
+      annos should contain(MarkAnnotation("~Top|HasMems>mem".rt, "Mem"))
+      annos should contain(MarkAnnotation("~Top|HasMems>syncReadMem".rt, "SyncReadMem"))
+    }
+    it("3.14: should not create memory ports") {
+      class Top() extends Module {
+        val i = Definition(new HasMems())
+        i.mem(0) := 100.U // should be illegal!
+      }
+      val failure = intercept[ChiselException] {
+        getFirrtlAndAnnos(new Top)
+      }
+      assert(
+        failure.getMessage ==
+          "Cannot create a memory port in a different module (Top) than where the memory is (HasMems)."
+      )
+    }
   }
   describe("4: toDefinition") {
     it("4.0: should work on modules") {

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

@@ -258,4 +258,10 @@ object Examples {
     val i10 = Instance(tpDef1)
     val i11 = Instance(tpDef1)
   }
+
+  @instantiable
+  class HasMems() extends Module {
+    @public val mem = Mem(8, UInt(32.W))
+    @public val syncReadMem = SyncReadMem(8, UInt(32.W))
+  }
 }

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

@@ -331,6 +331,17 @@ class InstanceSpec extends ChiselFunSpec with Utils {
         @public override final lazy val y: Int = 4
       }
     }
+    it("3.13: should work with Mems/SyncReadMems") {
+      class Top() extends Module {
+        val i = Instance(Definition(new HasMems()))
+        mark(i.mem, "Mem")
+        mark(i.syncReadMem, "SyncReadMem")
+      }
+      val (_, annos) = getFirrtlAndAnnos(new Top)
+      annos.foreach { x => println(x.serialize) }
+      annos should contain(MarkAnnotation("~Top|Top/i:HasMems>mem".rt, "Mem"))
+      annos should contain(MarkAnnotation("~Top|Top/i:HasMems>syncReadMem".rt, "SyncReadMem"))
+    }
   }
   describe("4: toInstance") {
     it("4.0: should work on modules") {