Hierarchy API: make Mems lookupable (#2404) (#2410)

(cherry picked from commit 2a985ac376698a2e6300fbee13001d82d3e13989)

Co-authored-by: Deborah Soung <debs@sifive.com>

4 files changed, 53 insertions, 4 deletions

diff --git a/core/src/main/scala/chisel3/Mem.scala b/core/src/main/scala/chisel3/Mem.scala
index 3f37308c..36984a3a 100644
--- a/core/src/main/scala/chisel3/Mem.scala
+++ b/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
diff --git a/core/src/main/scala/chisel3/Module.scala b/core/src/main/scala/chisel3/Module.scala
index 3611f5dd..ed323504 100644
--- a/core/src/main/scala/chisel3/Module.scala
+++ b/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
diff --git a/core/src/main/scala/chisel3/experimental/hierarchy/Lookupable.scala b/core/src/main/scala/chisel3/experimental/hierarchy/Lookupable.scala
index a0c2209b..bc94f95d 100644
--- a/core/src/main/scala/chisel3/experimental/hierarchy/Lookupable.scala
+++ b/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,
diff --git a/core/src/main/scala/chisel3/internal/Builder.scala b/core/src/main/scala/chisel3/internal/Builder.scala
index 247be57a..1ffe54ab 100644
--- a/core/src/main/scala/chisel3/internal/Builder.scala
+++ b/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}")
       }