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

(cherry picked from commit 2a985ac376698a2e6300fbee13001d82d3e13989)

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

1 files changed, 40 insertions, 1 deletions

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,