From 0531cb53d3cedaff33c2a280e34418f6af5bc6a1 Mon Sep 17 00:00:00 2001
From: Jack Koenig <koenig@sifive.com>
Date: Tue, 29 Jun 2021 15:34:18 -0700
Subject: [PATCH] Restore aop.Select behavior for CloneModuleAsRecord

---
 core/src/main/scala/chisel3/Module.scala      |  4 ++-
 src/main/scala/chisel3/aop/Select.scala       |  6 ++++-
 .../scala/chiselTests/aop/SelectSpec.scala    | 25 +++++++++++++++++++
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/core/src/main/scala/chisel3/Module.scala b/core/src/main/scala/chisel3/Module.scala
index 8a914fbcc5d..f82b6a7bb8b 100644
--- a/core/src/main/scala/chisel3/Module.scala
+++ b/core/src/main/scala/chisel3/Module.scala
@@ -184,12 +184,14 @@ package internal {
 
   object BaseModule {
     // Private internal class to serve as a _parent for Data in cloned ports
-    private[chisel3] class ModuleClone(_proto: BaseModule) extends BaseModule {
+    private[chisel3] class ModuleClone(private[chisel3] val _proto: BaseModule) extends BaseModule {
       // ClonePorts that hold the bound ports for this module
       // Used for setting the refs of both this module and the Record
       private[BaseModule] var _portsRecord: Record = _
       // Don't generate a component, but point to the one for the cloned Module
       private[chisel3] def generateComponent(): Option[Component] = {
+        require(!_closed, "Can't generate module more than once")
+        _closed = true
         _component = _proto._component
         None
       }
diff --git a/src/main/scala/chisel3/aop/Select.scala b/src/main/scala/chisel3/aop/Select.scala
index b9ad808b51c..08cf40fff7c 100644
--- a/src/main/scala/chisel3/aop/Select.scala
+++ b/src/main/scala/chisel3/aop/Select.scala
@@ -6,6 +6,7 @@ import chisel3._
 import chisel3.experimental.{BaseModule, FixedPoint}
 import chisel3.internal.HasId
 import chisel3.internal.firrtl._
+import chisel3.internal.BaseModule.ModuleClone
 import firrtl.annotations.ReferenceTarget
 
 import scala.collection.mutable
@@ -82,7 +83,10 @@ object Select {
     check(module)
     module._component.get match {
       case d: DefModule => d.commands.collect {
-        case i: DefInstance => i.id
+        case i: DefInstance => i.id match {
+          case clone: ModuleClone => clone._proto
+          case other              => other
+        }
       }
       case other => Nil
     }
diff --git a/src/test/scala/chiselTests/aop/SelectSpec.scala b/src/test/scala/chiselTests/aop/SelectSpec.scala
index 91353f5a601..d34f4391589 100644
--- a/src/test/scala/chiselTests/aop/SelectSpec.scala
+++ b/src/test/scala/chiselTests/aop/SelectSpec.scala
@@ -153,5 +153,30 @@ class SelectSpec extends ChiselFlatSpec {
     assert(bbs.size == 1)
   }
 
+  "CloneModuleAsRecord" should "show up in Select aspects as duplicates" in {
+    import chisel3.experimental.CloneModuleAsRecord
+    class Child extends RawModule {
+      val in = IO(Input(UInt(8.W)))
+      val out = IO(Output(UInt(8.W)))
+      out := in
+    }
+    class Top extends MultiIOModule {
+      val in = IO(Input(UInt(8.W)))
+      val out = IO(Output(UInt(8.W)))
+      val inst0 = Module(new Child)
+      val inst1 = CloneModuleAsRecord(inst0)
+      inst0.in := in
+      inst1("in") := inst0.out
+      out := inst1("out")
+    }
+    val top = ChiselGeneratorAnnotation(() => {
+      new Top()
+    }).elaborate
+      .collectFirst { case DesignAnnotation(design: Top) => design }
+      .get
+    val mods = Select.collectDeep(top) { case mod => mod }
+    mods should equal (Seq(top, top.inst0, top.inst0))
+  }
+
 }