Use Data refs for name prefixing with aggregate elements (#1616)

* Use Data refs for name prefixing with aggregate elements

Vecs set the refs of their elements upon construction of those elements.
In the past, Records haven't set their elements refs until module close,
but it can be done sooner. Doing it upon binding means that refs will at
least be available for Records used in hardware elements. Since only
bound Data can be connected to anyway, Aggregate elements being
connected to will always have a ref which we can then use for creating
naming prefixes.

* Add tighter correctness checks

* Handle more cases in connection prefixing

Add support for forcing setRef to override a previous setting. This
is only used by BlackBox ports which need to drop their io prefix.

Also add a Try() around Data.bindingToString which sometimes throws
exceptions when being used to .toString a Data in an error message.

* Strip trailing spaces in names in compiler plugin

7 files changed, 95 insertions, 22 deletions

diff --git a/core/src/main/scala/chisel3/Aggregate.scala b/core/src/main/scala/chisel3/Aggregate.scala
index 9ccf7dbb..6c11b2db 100644
--- a/core/src/main/scala/chisel3/Aggregate.scala
+++ b/core/src/main/scala/chisel3/Aggregate.scala
@@ -479,9 +479,21 @@ trait VecLike[T <: Data] extends collection.IndexedSeq[T] with HasId with Source
   * RTL writers should use [[Bundle]].  See [[Record#elements]] for an example.
   */
 abstract class Record(private[chisel3] implicit val compileOptions: CompileOptions) extends Aggregate {
+
+  // Doing this earlier than onModuleClose allows field names to be available for prefixing the names
+  // of hardware created when connecting to one of these elements
+  private def setElementRefs(): Unit = {
+    // Since elements is a map, it is impossible for two elements to have the same
+    // identifier; however, Namespace sanitizes identifiers to make them legal for Firrtl/Verilog
+    // which can cause collisions
+    val _namespace = Namespace.empty
+    for ((name, elt) <- elements) { elt.setRef(this, _namespace.name(name, leadingDigitOk=true)) }
+  }
+
   private[chisel3] override def bind(target: Binding, parentDirection: SpecifiedDirection): Unit = {
     try {
       super.bind(target, parentDirection)
+      setElementRefs()
     } catch {  // nasty compatibility mode shim, where anything flies
       case e: MixedDirectionAggregateException if !compileOptions.dontAssumeDirectionality =>
         val resolvedDirection = SpecifiedDirection.fromParent(parentDirection, specifiedDirection)
@@ -633,13 +645,11 @@ abstract class Record(private[chisel3] implicit val compileOptions: CompileOptio
     case _ => false
   }
 
-  // NOTE: This sets up dependent references, it can be done before closing the Module
   private[chisel3] override def _onModuleClose: Unit = {
-    // Since Bundle names this via reflection, it is impossible for two elements to have the same
-    // identifier; however, Namespace sanitizes identifiers to make them legal for Firrtl/Verilog
-    // which can cause collisions
-    val _namespace = Namespace.empty
-    for ((name, elt) <- elements) { elt.setRef(this, _namespace.name(name, leadingDigitOk=true)) }
+    // This is usually done during binding, but these must still be set for unbound Records
+    if (this.binding.isEmpty) {
+      setElementRefs()
+    }
   }
 
   private[chisel3] final def allElements: Seq[Element] = elements.toIndexedSeq.flatMap(_._2.allElements)
diff --git a/core/src/main/scala/chisel3/BlackBox.scala b/core/src/main/scala/chisel3/BlackBox.scala
index b204e301..1b093ee1 100644
--- a/core/src/main/scala/chisel3/BlackBox.scala
+++ b/core/src/main/scala/chisel3/BlackBox.scala
@@ -165,7 +165,8 @@ abstract class BlackBox(val params: Map[String, Param] = Map.empty[String, Param
     // Long term solution will be to define BlackBox IO differently as part of
     //   it not descending from the (current) Module
     for ((name, port) <- namedPorts) {
-      port.setRef(ModuleIO(this, _namespace.name(name)))
+      // We have to force override the _ref because it was set during IO binding
+      port.setRef(ModuleIO(this, _namespace.name(name)), force = true)
     }
 
     // We need to call forceName and onModuleClose on all of the sub-elements
diff --git a/core/src/main/scala/chisel3/Data.scala b/core/src/main/scala/chisel3/Data.scala
index bb45b4d8..5d398aa6 100644
--- a/core/src/main/scala/chisel3/Data.scala
+++ b/core/src/main/scala/chisel3/Data.scala
@@ -9,6 +9,8 @@ import chisel3.internal._
 import chisel3.internal.firrtl._
 import chisel3.internal.sourceinfo.{DeprecatedSourceInfo, SourceInfo, SourceInfoTransform, UnlocatableSourceInfo}
 
+import scala.util.Try
+
 /** User-specified directions.
   */
 sealed abstract class SpecifiedDirection
@@ -374,7 +376,7 @@ abstract class Data extends HasId with NamedComponent with SourceInfoDoc {
   // Provides a unhelpful fallback for literals, which should have custom rendering per
   // Data-subtype.
   // TODO Is this okay for sample_element? It *shouldn't* be visible to users
-  protected def bindingToString: String = topBindingOpt match {
+  protected def bindingToString: String = Try(topBindingOpt match {
     case None => ""
     case Some(OpBinding(enclosure, _)) => s"(OpResult in ${enclosure.desiredName})"
     case Some(MemoryPortBinding(enclosure, _)) => s"(MemPort in ${enclosure.desiredName})"
@@ -389,7 +391,7 @@ abstract class Data extends HasId with NamedComponent with SourceInfoDoc {
     case Some(DontCareBinding()) => s"(DontCare)"
     case Some(ElementLitBinding(litArg)) => s"(unhandled literal)"
     case Some(BundleLitBinding(litMap)) => s"(unhandled bundle literal)"
-  }
+  }).getOrElse("")
 
   // Return ALL elements at root of this type.
   // Contasts with flatten, which returns just Bits
diff --git a/core/src/main/scala/chisel3/internal/Builder.scala b/core/src/main/scala/chisel3/internal/Builder.scala
index ffa99913..3988ac68 100644
--- a/core/src/main/scala/chisel3/internal/Builder.scala
+++ b/core/src/main/scala/chisel3/internal/Builder.scala
@@ -198,7 +198,10 @@ private[chisel3] trait HasId extends InstanceId {
     }
 
   private var _ref: Option[Arg] = None
-  private[chisel3] def setRef(imm: Arg): Unit = _ref = Some(imm)
+  private[chisel3] def setRef(imm: Arg, force: Boolean = false): Unit = {
+    assert(force || _ref.isEmpty, s"Internal Error, setRef for $this called twice! first ${_ref.get}, second $imm")
+    _ref = Some(imm)
+  }
   private[chisel3] def setRef(parent: HasId, name: String): Unit = setRef(Slot(Node(parent), name))
   private[chisel3] def setRef(parent: HasId, index: Int): Unit = setRef(Index(Node(parent), ILit(index)))
   private[chisel3] def setRef(parent: HasId, index: UInt): Unit = setRef(Index(Node(parent), index.ref))
@@ -369,13 +372,12 @@ private[chisel3] object Builder {
     */
   def pushPrefix(d: HasId): Boolean = {
     def buildAggName(id: HasId): Option[String] = {
-      // TODO This is slow, can we store this information upon binding?
-      def getSubName(field: Data, parent: Data): Option[String] = parent match {
-        case vec: Vec[_] =>
-          val idx = vec.indexOf(field)
-          if (idx >= 0) Some(idx.toString) else None
-        case rec: Record => rec.elements.collectFirst { case (name, d) if d == field => name }
-        case _ => Builder.exception(s"Shouldn't see non-Aggregate $parent as parent in ChildBinding!")
+      def getSubName(field: Data): Option[String] = field.getOptionRef.flatMap {
+        case Slot(_, field) => Some(field) // Record
+        case Index(_, ILit(n)) => Some(n.toString) // Vec static indexing
+        case Index(_, ULit(n, _)) => Some(n.toString) // Vec lit indexing
+        case Index(_, _: Node) => None // Vec dynamic indexing
+        case ModuleIO(_, n) => Some(n) // BlackBox port
       }
       def map2[A, B](a: Option[A], b: Option[A])(f: (A, A) => B): Option[B] =
         a.flatMap(ax => b.map(f(ax, _)))
@@ -384,7 +386,7 @@ private[chisel3] object Builder {
         case (_: WireBinding | _: RegBinding | _: MemoryPortBinding | _: OpBinding) => data.seedOpt
         case ChildBinding(parent) => recData(parent).map { p =>
           // And name of the field if we have one, we don't for dynamic indexing of Vecs
-          getSubName(data, parent).map(p + "_" + _).getOrElse(p)
+          getSubName(data).map(p + "_" + _).getOrElse(p)
         }
         case SampleElementBinding(parent) => recData(parent)
         case PortBinding(mod) if Builder.currentModule.contains(mod) => data.seedOpt
diff --git a/plugin/src/main/scala-2.12/chisel3/internal/plugin/ChiselPlugin.scala b/plugin/src/main/scala-2.12/chisel3/internal/plugin/ChiselPlugin.scala
index 6e4f3ece..9f90ba16 100644
--- a/plugin/src/main/scala-2.12/chisel3/internal/plugin/ChiselPlugin.scala
+++ b/plugin/src/main/scala-2.12/chisel3/internal/plugin/ChiselPlugin.scala
@@ -124,27 +124,30 @@ class ChiselComponent(val global: Global) extends PluginComponent with TypingTra
       dd.symbol.logicallyEnclosingMember.thisType <:< inferType(tq"chisel3.Bundle")
     }
 
+    private def stringFromTermName(name: TermName): String =
+      name.toString.trim() // Remove trailing space (Scalac implementation detail)
+
     // Method called by the compiler to modify source tree
     override def transform(tree: Tree): Tree = tree match {
       // Check if a subtree is a candidate
       case dd @ ValDef(mods, name, tpt, rhs) if okVal(dd) =>
       // If a Data and in a Bundle, just get the name but not a prefix
         if (shouldMatchData(dd) && inBundle(dd)) {
-          val TermName(str: String) = name
+          val str = stringFromTermName(name)
           val newRHS = transform(rhs)  // chisel3.internal.plugin.autoNameRecursively
           val named = q"chisel3.internal.plugin.autoNameRecursively($str, $newRHS)"
           treeCopy.ValDef(dd, mods, name, tpt, localTyper typed named)
         }
         // If a Data or a Memory, get the name and a prefix
         else if (shouldMatchDataOrMem(dd)) {
-          val TermName(str: String) = name
+          val str = stringFromTermName(name)
           val newRHS = transform(rhs)
           val prefixed = q"chisel3.experimental.prefix.apply[$tpt](name=$str)(f=$newRHS)"
           val named = q"chisel3.internal.plugin.autoNameRecursively($str, $prefixed)"
           treeCopy.ValDef(dd, mods, name, tpt, localTyper typed named)
         // If an instance, just get a name but no prefix
         } else if (shouldMatchModule(dd)) {
-          val TermName(str: String) = name
+          val str = stringFromTermName(name)
           val newRHS = transform(rhs)
           val named = q"chisel3.internal.plugin.autoNameRecursively($str, $newRHS)"
           treeCopy.ValDef(dd, mods, name, tpt, localTyper typed named)
diff --git a/src/test/scala/chiselTests/Mem.scala b/src/test/scala/chiselTests/Mem.scala
index b66ec42f..8bcd3aac 100644
--- a/src/test/scala/chiselTests/Mem.scala
+++ b/src/test/scala/chiselTests/Mem.scala
@@ -96,6 +96,51 @@ class HugeCMemTester(size: BigInt) extends BasicTester {
   }
 }
 
+class SyncReadMemBundleTester extends BasicTester {
+  val (cnt, _) = Counter(true.B, 5)
+  val tpe = new Bundle {
+    val foo = UInt(2.W)
+  }
+  val mem = SyncReadMem(2, tpe)
+  val rdata = mem.read(cnt - 1.U, cnt =/= 0.U)
+
+  switch (cnt) {
+    is (0.U) {
+      val w = Wire(tpe)
+      w.foo := 3.U
+      mem.write(cnt, w)
+    }
+    is (1.U) {
+      val w = Wire(tpe)
+      w.foo := 2.U
+      mem.write(cnt, w)
+    }
+    is (2.U) { assert(rdata.foo === 3.U) }
+    is (3.U) { assert(rdata.foo === 2.U) }
+    is (4.U) { stop() }
+  }
+}
+
+class MemBundleTester extends BasicTester {
+  val tpe = new Bundle {
+    val foo = UInt(2.W)
+  }
+  val mem = Mem(2, tpe)
+
+  // Circuit style tester is definitely the wrong abstraction here
+  val (cnt, wrap) = Counter(true.B, 2)
+  mem(0) := {
+    val w = Wire(tpe)
+    w.foo := 1.U
+    w
+  }
+
+  when (cnt === 1.U) {
+    assert(mem.read(0.U).foo === 1.U)
+    stop()
+  }
+}
+
 class MemorySpec extends ChiselPropSpec {
   property("Mem of Vec should work") {
     assertTesterPasses { new MemVecTester }
@@ -105,6 +150,14 @@ class MemorySpec extends ChiselPropSpec {
     assertTesterPasses { new SyncReadMemTester }
   }
 
+  property("SyncReadMems of Bundles should work") {
+    assertTesterPasses { new SyncReadMemBundleTester }
+  }
+
+  property("Mems of Bundles should work") {
+    assertTesterPasses { new MemBundleTester }
+  }
+
   property("SyncReadMem write collision behaviors should work") {
     assertTesterPasses { new SyncReadMemWriteCollisionTester }
   }
diff --git a/src/test/scala/chiselTests/naming/PrefixSpec.scala b/src/test/scala/chiselTests/naming/PrefixSpec.scala
index 27a9fd39..83408dea 100644
--- a/src/test/scala/chiselTests/naming/PrefixSpec.scala
+++ b/src/test/scala/chiselTests/naming/PrefixSpec.scala
@@ -285,6 +285,7 @@ class PrefixSpec extends ChiselPropSpec with Utils {
         wire.y := RegNext(3.U)
         wire.vec(0) := RegNext(3.U)
         wire.vec(wire.x) := RegNext(3.U)
+        wire.vec(1.U) := RegNext(3.U)
       }
     }
     aspectTest(() => new Test) {
@@ -293,7 +294,8 @@ class PrefixSpec extends ChiselPropSpec with Utils {
           "wire_x_REG",
           "wire_y_REG",
           "wire_vec_0_REG",
-          "wire_vec_REG"
+          "wire_vec_REG",
+          "wire_vec_1_REG"
         ))
     }
   }