Improve reference & name resolution in ReplSeqMem (#352)

* fix imports in InferReadWrite

* improve reference & name resolution in ReplSeqMem

* add comments

5 files changed, 38 insertions, 23 deletions

diff --git a/src/main/scala/firrtl/passes/memlib/InferReadWrite.scala b/src/main/scala/firrtl/passes/memlib/InferReadWrite.scala
index 9adbdd95..ffdea1f2 100644
--- a/src/main/scala/firrtl/passes/memlib/InferReadWrite.scala
+++ b/src/main/scala/firrtl/passes/memlib/InferReadWrite.scala
@@ -32,9 +32,8 @@ import firrtl.ir._
 import firrtl.Mappers._
 import firrtl.PrimOps._
 import firrtl.Utils.{one, zero, BoolType}
-import firrtl.passes.memlib._
 import MemPortUtils.memPortField
-import AnalysisUtils.{Connects, getConnects, getOrigin}
+import firrtl.passes.memlib.AnalysisUtils.{Connects, getConnects, getOrigin}
 import WrappedExpression.weq
 import Annotations._
 
diff --git a/src/main/scala/firrtl/passes/memlib/MemIR.scala b/src/main/scala/firrtl/passes/memlib/MemIR.scala
index 6dca5961..54441481 100644
--- a/src/main/scala/firrtl/passes/memlib/MemIR.scala
+++ b/src/main/scala/firrtl/passes/memlib/MemIR.scala
@@ -19,7 +19,7 @@ case class DefAnnotatedMemory(
     readwriters: Seq[String],
     readUnderWrite: Option[String],
     maskGran: Option[BigInt],
-    memRef: Option[String]
+    memRef: Option[(String, String)] /* (Module, Mem) */
     //pins: Seq[Pin],
     ) extends Statement with IsDeclaration {
   def serialize: String = this.toMem.serialize
diff --git a/src/main/scala/firrtl/passes/memlib/RenameAnnotatedMemoryPorts.scala b/src/main/scala/firrtl/passes/memlib/RenameAnnotatedMemoryPorts.scala
index 168a6a48..81242810 100644
--- a/src/main/scala/firrtl/passes/memlib/RenameAnnotatedMemoryPorts.scala
+++ b/src/main/scala/firrtl/passes/memlib/RenameAnnotatedMemoryPorts.scala
@@ -35,8 +35,7 @@ object RenameAnnotatedMemoryPorts extends Pass {
    *  E.g.:
    *    - ("m.read.addr") becomes (m.R0.addr)
    */
-  def getMemPortMap(m: DefAnnotatedMemory): MemPortMap = {
-    val memPortMap = new MemPortMap
+  def getMemPortMap(m: DefAnnotatedMemory, memPortMap: MemPortMap) {
     val defaultFields = Seq("addr", "en", "clk")
     val rFields = defaultFields :+ "data"
     val wFields = rFields :+ "mask"
@@ -51,7 +50,6 @@ object RenameAnnotatedMemoryPorts extends Pass {
     updateMemPortMap(m.readers, rFields, "R")
     updateMemPortMap(m.writers, wFields, "W")
     updateMemPortMap(m.readwriters, rwFields, "RW")
-    memPortMap
   }
 
   /** Replaces candidate memories with memories with standard port names
@@ -60,7 +58,7 @@ object RenameAnnotatedMemoryPorts extends Pass {
   def updateMemStmts(memPortMap: MemPortMap)(s: Statement): Statement = s match {
     case m: DefAnnotatedMemory =>
       val updatedMem = createMemProto(m)
-      memPortMap ++= getMemPortMap(m)
+      getMemPortMap(m, memPortMap)
       updatedMem
     case s => s map updateMemStmts(memPortMap)
   }
diff --git a/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala b/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala
index 91e4efea..a52f7d38 100644
--- a/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala
+++ b/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala
@@ -148,7 +148,23 @@ class ReplaceMemMacros(writer: ConfWriter) extends Pass {
     }
   }
 
+  /** 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 {
+      case m: DefAnnotatedMemory => m.memRef match {
+        case None => nameMap(mname -> m.name) = namespace newName m.name
+        case Some(_) =>
+      }
+      case _ =>
+    }
+    s map constructNameMap(namespace, nameMap, mname)
+  }
+
   def updateMemStmts(namespace: Namespace,
+                     nameMap: NameMap,
+                     mname: String,
                      memPortMap: MemPortMap,
                      memMods: Modules)
                      (s: Statement): Statement = s match {
@@ -160,28 +176,30 @@ class ReplaceMemMacros(writer: ConfWriter) extends Pass {
       m.memRef match {
         case None =>
           // prototype mem
-          val newWrapperName = namespace newName m.name
-          val newMemBBName = namespace newName s"${m.name}_ext"
+          val newWrapperName = nameMap(mname -> m.name)
+          val newMemBBName = namespace newName s"${newWrapperName}_ext"
           val newMem = m copy (name = newMemBBName)
           memMods ++= createMemModule(newMem, newWrapperName)
-          WDefInstance(m.info, m.name, newWrapperName, UnknownType) 
-        case Some(ref: String) =>
-          WDefInstance(m.info, m.name, ref, UnknownType) 
+          WDefInstance(m.info, m.name, newWrapperName, UnknownType)
+        case Some((module, mem)) =>
+          WDefInstance(m.info, m.name, nameMap(module -> mem), UnknownType)
       }
-    case sx => sx map updateMemStmts(namespace, memPortMap, memMods)
+    case sx => sx map updateMemStmts(namespace, nameMap, mname, memPortMap, memMods)
   }
 
-  def updateMemMods(namespace: Namespace, memMods: Modules)(m: DefModule) = {
+  def updateMemMods(namespace: Namespace, nameMap: NameMap, memMods: Modules)(m: DefModule) = {
     val memPortMap = new MemPortMap
 
-    (m map updateMemStmts(namespace, memPortMap, memMods)
+    (m map updateMemStmts(namespace, nameMap, m.name, memPortMap, memMods)
        map updateStmtRefs(memPortMap))
   }
 
   def run(c: Circuit) = {
     val namespace = Namespace(c)
     val memMods = new Modules
-    val modules = c.modules map updateMemMods(namespace, memMods)
+    val nameMap = new NameMap
+    c.modules map (m => m map constructNameMap(namespace, nameMap, m.name))
+    val modules = c.modules map updateMemMods(namespace, nameMap, memMods)
     // print conf
     writer.serialize()
     c copy (modules = modules ++ memMods)
diff --git a/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala b/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala
index 108f6fd4..2112ca27 100644
--- a/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala
+++ b/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala
@@ -13,24 +13,24 @@ object ResolveMemoryReference extends Pass {
 
   def name = "Resolve Memory Reference"
 
-  type AnnotatedMemories = collection.mutable.ArrayBuffer[DefAnnotatedMemory]
+  type AnnotatedMemories = collection.mutable.ArrayBuffer[(String, DefAnnotatedMemory)]
 
   /** If a candidate memory is identical except for name to another, add an
     *   annotation that references the name of the other memory.
     */
-  def updateMemStmts(uniqueMems: AnnotatedMemories)(s: Statement): Statement = s match {
+  def updateMemStmts(mname: String, uniqueMems: AnnotatedMemories)(s: Statement): Statement = s match {
     case m: DefAnnotatedMemory =>
-      uniqueMems find (x => eqMems(x, m)) match {
+      uniqueMems find (x => eqMems(x._2, m)) match {
         case None =>
-          uniqueMems += m
+          uniqueMems += (mname -> m)
           m
-        case Some(proto) => m copy (memRef = Some(proto.name))
+        case Some((module, proto)) => m copy (memRef = Some(module -> proto.name))
       }
-    case s => s map updateMemStmts(uniqueMems)
+    case s => s map updateMemStmts(mname, uniqueMems)
   }
 
   def run(c: Circuit) = {
     val uniqueMems = new AnnotatedMemories
-    c copy (modules = c.modules map (_ map updateMemStmts(uniqueMems)))
+    c copy (modules = c.modules map (m => m map updateMemStmts(m.name, uniqueMems)))
   }
 }