Make nested IsInstantiables with Data in them work (#2761) (#2766)

* Add unit test for Issue 2760

* checkpoint: Fix for nested instance

* remove comments about stuff not working

* make the test check the output a little more

* relax the requirement on returning empty ioMap

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

* Update core/src/main/scala/chisel3/Data.scala

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

Co-authored-by: Jack Koenig <koenig@sifive.com>

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

Co-authored-by: Jack Koenig <koenig@sifive.com>

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

* Add another unit test which unfortunately still passes

* Update core/src/main/scala/chisel3/Data.scala

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

Co-authored-by: Jack Koenig <koenig@sifive.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
(cherry picked from commit 1f9f26dc2bffcb4cc4daf2dc16c5cb455c6769ef)

Co-authored-by: Megan Wachs <megan@sifive.com>

3 files changed, 89 insertions, 4 deletions

diff --git a/core/src/main/scala/chisel3/experimental/hierarchy/Lookupable.scala b/core/src/main/scala/chisel3/experimental/hierarchy/Lookupable.scala
index c83479b0..aa35455d 100644
--- a/core/src/main/scala/chisel3/experimental/hierarchy/Lookupable.scala
+++ b/core/src/main/scala/chisel3/experimental/hierarchy/Lookupable.scala
@@ -335,11 +335,20 @@ object Lookupable {
       }
       def instanceLookup[A](that: A => B, instance: Instance[A]): C = {
         val ret = that(instance.proto)
-        val ioMap: Option[Map[Data, Data]] = instance.underlying match {
-          case Clone(x: ModuleClone[_]) => Some(x.ioMap)
-          case Proto(x: BaseModule) => Some(x.getChiselPorts.map { case (_, data) => data -> data }.toMap)
-          case _ => None
+
+        def getIoMap(hierarchy: Hierarchy[_]): Option[Map[Data, Data]] = {
+          hierarchy.underlying match {
+            case Clone(x: ModuleClone[_]) => Some(x.ioMap)
+            case Proto(x: BaseModule) => Some(x.getChiselPorts.map { case (_, data) => data -> data }.toMap)
+            case Clone(x: InstantiableClone[_]) => getIoMap(x._innerContext)
+            case Clone(x: InstanceClone[_]) => None
+            case other => {
+              Builder.exception(s"Internal Error! Unexpected case where we can't get IO Map: $other")
+            }
+          }
         }
+        val ioMap = getIoMap(instance)
+
         if (isView(ret)) {
           cloneViewToContext(ret, instance.cache, ioMap, instance.getInnerDataContext)
         } else {
diff --git a/src/test/scala/chiselTests/experimental/hierarchy/Examples.scala b/src/test/scala/chiselTests/experimental/hierarchy/Examples.scala
index fa26cbde..27725c49 100644
--- a/src/test/scala/chiselTests/experimental/hierarchy/Examples.scala
+++ b/src/test/scala/chiselTests/experimental/hierarchy/Examples.scala
@@ -271,4 +271,70 @@ object Examples {
     @public val mem = Mem(8, UInt(32.W))
     @public val syncReadMem = SyncReadMem(8, UInt(32.W))
   }
+
+  @instantiable
+  class LeafInstantiable(val bundle: Data) {
+    @public val bundle = bundle
+  }
+
+  @instantiable
+  class NestedInstantiable(val in: LeafInstantiable, val out: LeafInstantiable) {
+    @public val in = in
+    @public val out = out
+  }
+
+  @instantiable
+  class AddOneNestedInstantiableData(width: Int) extends Module {
+    @public val in = IO(Input(UInt(width.W)))
+    @public val out = IO(Output(UInt(width.W)))
+    out := in + 1.U
+
+    @public val leafOut = new LeafInstantiable(out)
+    @public val leafIn = new LeafInstantiable(in)
+    @public val nested = new NestedInstantiable(in = leafIn, out = leafOut)
+
+  }
+
+  class AddTwoNestedInstantiableData(width: Int) extends Module {
+    val in = IO(Input(UInt(width.W)))
+    val out = IO(Output(UInt(width.W)))
+    val addOneDef = Definition(new AddOneNestedInstantiableData(width))
+    val i0 = Instance(addOneDef)
+    val i1 = Instance(addOneDef)
+    i0.in := in
+    i1.in := i0.out
+    out := i1.out
+
+    // both are equivalent to the above
+    i1.leafIn.bundle := i0.leafOut.bundle
+    i1.nested.in.bundle := i0.nested.out.bundle
+  }
+
+  class AddTwoNestedInstantiableDataSubmodule(addOneDef: Definition[AddOneNestedInstantiableData]) extends Module {
+    val in = IO(Input(UInt(addOneDef.in.getWidth.W)))
+    val out = IO(Output(UInt(addOneDef.out.getWidth.W)))
+    val i0 = Instance(addOneDef)
+    val i1 = Instance(addOneDef)
+    i0.in := in
+    i1.in := i0.out
+    out := i1.out
+
+    // both are equivalent to the above
+    i1.leafIn.bundle := i0.leafOut.bundle
+    i1.nested.in.bundle := i0.nested.out.bundle
+  }
+
+  class AddTwoNestedInstantiableDataWrapper(width: Int) extends Module {
+    val in = IO(Input(UInt(width.W)))
+    val out = IO(Output(UInt(width.W)))
+
+    val original = Module(new AddOneNestedInstantiableData(width))
+    val copy = Module(new AddTwoNestedInstantiableDataSubmodule(original.toDefinition))
+
+    original.in := in
+    copy.in := original.out
+    out := copy.out
+
+  }
+
 }
diff --git a/src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala b/src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala
index 8d8f7ea5..6596cd51 100644
--- a/src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala
+++ b/src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala
@@ -364,6 +364,16 @@ class InstanceSpec extends ChiselFunSpec with Utils {
       annos should contain(MarkAnnotation("~Top|Top/i:HasMems>mem".rt, "Mem"))
       annos should contain(MarkAnnotation("~Top|Top/i:HasMems>syncReadMem".rt, "SyncReadMem"))
     }
+    it("(3.p): should make connectable IOs on nested IsInstantiables that have IO Datas in them") {
+      val (chirrtl, _) = getFirrtlAndAnnos(new AddTwoNestedInstantiableData(4))
+      exactly(3, chirrtl.serialize.split('\n')) should include("i1.in <= i0.out")
+    }
+    it(
+      "(3.q): should make connectable IOs on nested IsInstantiables's Data when the Instance and Definition do not have the same parent"
+    ) {
+      val (chirrtl, _) = getFirrtlAndAnnos(new AddTwoNestedInstantiableDataWrapper(4))
+      exactly(3, chirrtl.serialize.split('\n')) should include("i1.in <= i0.out")
+    }
   }
   describe("4: toInstance") {
     it("4.0: should work on modules") {