Repl seq mem renaming (#1286)

* Consume NoDedupMemAnnotations in ResolveMemoryReference

The ComponentName being pointed to by the annotation no longer exists
after ReplaceSeqMems so we should consume the annotations

* Support renaming in ReplaceMemMacros

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

2 files changed, 68 insertions, 11 deletions

diff --git a/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala b/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala
index f81dc71b..7b3608a3 100644
--- a/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala
+++ b/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala
@@ -12,12 +12,25 @@ import MemTransformUtils._
 import firrtl.annotations._
 import wiring._
 
+import scala.collection.mutable
+
 
 /** Annotates the name of the pins to add for WiringTransform */
 case class PinAnnotation(pins: Seq[String]) extends NoTargetAnnotation
 
 object ReplaceMemMacros {
   class UnsupportedBlackboxMemoryException(msg: String) extends PassException(msg)
+
+  /** Mapping from (module, memory name) pairs to wrapper Module names */
+  private type NameMap = mutable.HashMap[(String, String), String]
+
+  /** Mutable datastructure representing mapping of smems to extracted blackboxes and their wrappers
+    *
+    * - nameMap: mapping from every (module, memory name) to mem blackbox wrapper Module name
+    * - bbMap: mapping from wrapper Module name to (blackbox instance, blackbox module)
+    * - For bbMap: instance and module name match in the code, but that could be changed
+    */
+  class MemMapping(val nameMap: NameMap, val bbMap: mutable.HashMap[String, (String, String)])
 }
 
 /** Replace DefAnnotatedMemory with memory blackbox + wrapper + conf file.
@@ -28,6 +41,8 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
   def inputForm = MidForm
   def outputForm = MidForm
 
+  import ReplaceMemMacros._
+
   /** Return true if mask granularity is per bit, false if per byte or unspecified
     */
   private def getFillWMask(mem: DefAnnotatedMemory) = mem.maskGran match {
@@ -197,8 +212,6 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
     }
   }
 
-  /** Mapping from (module, memory name) pairs to blackbox names */
-  private type NameMap = collection.mutable.HashMap[(String, String), String]
   /** Construct NameMap by assigning unique names for each memory blackbox */
   def constructNameMap(namespace: Namespace, nameMap: NameMap, mname: String)(s: Statement): Statement = {
     s match {
@@ -211,12 +224,23 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
     s map constructNameMap(namespace, nameMap, mname)
   }
 
+  // For 1.2.x backwards compatibility
   def updateMemStmts(namespace: Namespace,
                      nameMap: NameMap,
                      mname: String,
                      memPortMap: MemPortMap,
-                     memMods: Modules)
-                     (s: Statement): Statement = s match {
+                     memMods: Modules
+                    )(s: Statement): Statement =
+    updateMemStmts(namespace, nameMap, mname, memPortMap, memMods, None)(s)
+
+  // memMapping is only Option for backwards compatibility
+  def updateMemStmts(namespace: Namespace,
+                     nameMap: NameMap,
+                     mname: String,
+                     memPortMap: MemPortMap,
+                     memMods: Modules,
+                     memMapping: Option[MemMapping]
+                    )(s: Statement): Statement = s match {
     case m: DefAnnotatedMemory =>
       if (m.maskGran.isEmpty) {
         m.writers foreach { w => memPortMap(s"${m.name}.$w.mask") = EmptyExpression }
@@ -228,18 +252,39 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
           val newWrapperName = nameMap(mname -> m.name)
           val newMemBBName = namespace newName s"${newWrapperName}_ext"
           val newMem = m copy (name = newMemBBName)
+          // Record for annotation renaming
+          memMapping.foreach { mapping =>
+            mapping.nameMap += ((mname, m.name) -> newWrapperName)
+            mapping.bbMap += newWrapperName -> (newMemBBName, newMemBBName)
+          }
           memMods ++= createMemModule(newMem, newWrapperName)
           WDefInstance(m.info, m.name, newWrapperName, UnknownType)
         case Some((module, mem)) =>
-          WDefInstance(m.info, m.name, nameMap(module -> mem), UnknownType)
+          val wrapperName = nameMap(module -> mem)
+          // Record for annotation renaming
+          memMapping.foreach(_.nameMap += ((mname, m.name) -> wrapperName))
+          WDefInstance(m.info, m.name, wrapperName, UnknownType)
       }
-    case sx => sx map updateMemStmts(namespace, nameMap, mname, memPortMap, memMods)
+    case sx => sx map updateMemStmts(namespace, nameMap, mname, memPortMap, memMods, memMapping)
   }
 
-  def updateMemMods(namespace: Namespace, nameMap: NameMap, memMods: Modules)(m: DefModule) = {
+
+  // For 1.2.x backwards compatibility
+  def updateMemMods(namespace: Namespace,
+                    nameMap: NameMap,
+                    memMods: Modules
+                   )(m: DefModule): DefModule =
+    updateMemMods(namespace, nameMap, memMods, None)(m)
+
+  // memMapping is only Option for backwards compatibility
+  def updateMemMods(namespace: Namespace,
+                    nameMap: NameMap,
+                    memMods: Modules,
+                    memMapping: Option[MemMapping]
+                   )(m: DefModule): DefModule = {
     val memPortMap = new MemPortMap
 
-    (m map updateMemStmts(namespace, nameMap, m.name, memPortMap, memMods)
+    (m map updateMemStmts(namespace, nameMap, m.name, memPortMap, memMods, memMapping)
        map updateStmtRefs(memPortMap))
   }
 
@@ -249,7 +294,18 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
     val memMods = new Modules
     val nameMap = new NameMap
     c.modules map (m => m map constructNameMap(namespace, nameMap, m.name))
-    val modules = c.modules map updateMemMods(namespace, nameMap, memMods)
+    val memMapping = new MemMapping(new NameMap, new mutable.HashMap)
+    val modules = c.modules map updateMemMods(namespace, nameMap, memMods, Some(memMapping))
+
+    // Rename replaced memories with new blackbox
+    val renames = RenameMap.create {
+      val top = CircuitTarget(c.main)
+      memMapping.nameMap.map { case ((mod, mem), wrap) =>
+        val (bbInst, bbMod) = memMapping.bbMap(wrap)
+        top.module(mod).ref(mem) -> Seq(top.module(mod).instOf(mem, wrap).instOf(bbInst, bbMod))
+      }
+    }
+
     // print conf
     writer.serialize()
     val pannos = state.annotations.collect { case a: PinAnnotation => a }
@@ -263,6 +319,6 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
         case m: ExtModule => SinkAnnotation(ModuleName(m.name, CircuitName(c.main)), pin)
       }
     } ++ state.annotations
-    CircuitState(c.copy(modules = modules ++ memMods), inputForm, annos)
+    CircuitState(c.copy(modules = modules ++ memMods), inputForm, annos, renames = Some(renames))
   }
 }
diff --git a/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala b/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala
index b0d3731f..094c2929 100644
--- a/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala
+++ b/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala
@@ -68,7 +68,8 @@ class ResolveMemoryReference extends Transform {
     val noDedups = state.annotations.collect {
       case NoDedupMemAnnotation(ComponentName(cn, ModuleName(mn, _))) => mn -> cn
     }
+    val annos = state.annotations.filterNot(_.isInstanceOf[NoDedupMemAnnotation])
     val noDedupMap: Map[String, Set[String]] = noDedups.groupBy(_._1).mapValues(_.map(_._2).toSet)
-    state.copy(circuit = run(state.circuit, noDedupMap))
+    state.copy(circuit = run(state.circuit, noDedupMap), annotations = annos)
   }
 }