Set refs for ModuleClone and ClonePorts in less hacky way

3 files changed, 39 insertions, 31 deletions

diff --git a/core/src/main/scala/chisel3/Module.scala b/core/src/main/scala/chisel3/Module.scala
index 77735583..8a914fbc 100644
--- a/core/src/main/scala/chisel3/Module.scala
+++ b/core/src/main/scala/chisel3/Module.scala
@@ -184,48 +184,44 @@ 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(_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] = {
-        _component = proto._component
+        _component = _proto._component
         None
       }
       // This module doesn't acutally exist in the FIRRTL so no initialization to do
       private[chisel3] def initializeInParent(parentCompileOptions: CompileOptions): Unit = ()
 
-      override def desiredName: String = proto.name
+      override def desiredName: String = _proto.name
+
+      private[chisel3] def setRefAndPortsRef(namespace: Namespace): Unit = {
+        val record = _portsRecord
+        // Use .forceName to re-use default name resolving behavior
+        record.forceName(None, default=this.desiredName, namespace)
+        // Now take the Ref that forceName set and convert it to the correct Arg
+        val instName = record.getRef match {
+          case Ref(name) => name
+          case bad => throwException(s"Internal Error! Cloned-module Record $record has unexpected ref $bad")
+        }
+        // Set both the record and the module to have the same instance name
+        record.setRef(ModuleCloneIO(_proto, instName), force=true) // force because we did .forceName first
+        this.setRef(Ref(instName))
+      }
     }
 
     /** Record type returned by CloneModuleAsRecord
       *
       * @note These are not true Data (the Record doesn't correspond to anything in the emitted
       * FIRRTL yet its elements *do*) so have some very specialized behavior.
-      * @param proto Optional pointer to the Module we are a clone of. Set for first instance, unset
-      *              for clones
       */
-    private[chisel3] class ClonePorts (proto: Option[BaseModule], elts: Data*)(implicit compileOptions: CompileOptions) extends Record {
+    private[chisel3] class ClonePorts (elts: Data*)(implicit compileOptions: CompileOptions) extends Record {
       val elements = ListMap(elts.map(d => d.instanceName -> d.cloneTypeFull): _*)
       def apply(field: String) = elements(field)
-      override def cloneType = (new ClonePorts(None, elts: _*)).asInstanceOf[this.type]
-
-      // Because ClonePorts instances are *not* created inside of their parent module, but rather,
-      // their parent's parent, we have to intercept the standard setRef and replace it with our own
-      // special Ref type.
-      // This only applies to ClonePorts created in cloneIORecord, any clones of these Records have
-      // normal behavior.
-      // Also, the name of ClonePorts Records needs to be propagated to their parent ModuleClone
-      // since we have no other way of setting the instance name for those.
-      private[chisel3] override def setRef(imm: Arg, force: Boolean): Unit = {
-        val immx = (proto, imm) match {
-          case (Some(mod), Ref(name)) =>
-            // Our _parent is a ModuleClone that needs its ref to match ours for .toAbsoluteTarget
-            _parent.foreach(_.setRef(Ref(name), force=true))
-            // Return a specialize-ref that will do the right thing
-            ModuleCloneIO(mod, name)
-          case _ => imm
-        }
-        super.setRef(immx, force)
-      }
+      override def cloneType = (new ClonePorts(elts: _*)).asInstanceOf[this.type]
     }
 
     // Recursively set the parent of the start Data and any children (eg. in an Aggregate)
@@ -243,12 +239,16 @@ package internal {
 
     private[chisel3] def cloneIORecord(proto: BaseModule)(implicit sourceInfo: SourceInfo, compileOptions: CompileOptions): ClonePorts = {
       require(proto.isClosed, "Can't clone a module before module close")
+      // Fake Module to serve as the _parent of the cloned ports
+      // We make this before clonePorts because we want it to come up first in naming in
+      // currentModule
+      val cloneParent = Module(new ModuleClone(proto))
       // We don't create this inside the ModuleClone because we need the ref to be set by the
       // currentModule (and not clonePorts)
-      val clonePorts = new ClonePorts(Some(proto), proto.getModulePorts: _*)
-      val cloneParent = Module(new ModuleClone(proto))
+      val clonePorts = new ClonePorts(proto.getModulePorts: _*)
       clonePorts.bind(PortBinding(cloneParent))
       setAllParents(clonePorts, Some(cloneParent))
+      cloneParent._portsRecord = clonePorts
       // Normally handled during Module construction but ClonePorts really lives in its parent's parent
       if (!compileOptions.explicitInvalidate) {
         pushCommand(DefInvalid(sourceInfo, clonePorts.ref))
diff --git a/core/src/main/scala/chisel3/RawModule.scala b/core/src/main/scala/chisel3/RawModule.scala
index 5bcd4dbd..4a60ca47 100644
--- a/core/src/main/scala/chisel3/RawModule.scala
+++ b/core/src/main/scala/chisel3/RawModule.scala
@@ -7,6 +7,7 @@ import scala.util.Try
 import scala.language.experimental.macros
 import chisel3.experimental.{BaseModule, BaseSim}
 import chisel3.internal._
+import chisel3.internal.BaseModule.ModuleClone
 import chisel3.internal.Builder._
 import chisel3.internal.firrtl._
 import chisel3.internal.sourceinfo.UnlocatableSourceInfo
@@ -74,6 +75,7 @@ abstract class RawModule(implicit moduleCompileOptions: CompileOptions)
     // All suggestions are in, force names to every node.
     for (id <- getIds) {
       id match {
+        case id: ModuleClone => id.setRefAndPortsRef(_namespace) // special handling
         case id: BaseModule => id.forceName(None, default=id.desiredName, _namespace)
         case id: MemBase[_] => id.forceName(None, default="MEM", _namespace)
         case id: BaseSim => id.forceName(None, default="SIM", _namespace)
diff --git a/core/src/main/scala/chisel3/internal/firrtl/IR.scala b/core/src/main/scala/chisel3/internal/firrtl/IR.scala
index 1d77802b..a4f6d26d 100644
--- a/core/src/main/scala/chisel3/internal/firrtl/IR.scala
+++ b/core/src/main/scala/chisel3/internal/firrtl/IR.scala
@@ -161,15 +161,21 @@ case class IntervalLit(n: BigInt, w: Width, binaryPoint: BinaryPoint) extends Li
 }
 
 case class Ref(name: String) extends Arg
+/** Arg for ports of Modules
+  * @param mod the module this port belongs to
+  * @param name the name of the port
+  */
 case class ModuleIO(mod: BaseModule, name: String) extends Arg {
   override def fullName(ctx: Component): String =
     if (mod eq ctx.id) name else s"${mod.getRef.name}.$name"
 }
-// For use with CloneModuleAsRecord
-// Note that `name` is the name of the module instance whereas in ModuleIO it's the name of the port
-// The names of ports inside of a ModuleCloneIO are the names of the Slots
+/** Ports of cloned modules (CloneModuleAsRecord)
+  * @param mod The original module for which these ports are a clone
+  * @param name the name of the module instance
+  */
 case class ModuleCloneIO(mod: BaseModule, name: String) extends Arg {
   override def fullName(ctx: Component): String =
+    // NOTE: mod eq ctx.id only occurs in Target and Named-related APIs
     if (mod eq ctx.id) "" else name
 }
 case class Slot(imm: Node, name: String) extends Arg {