Fix NoDedupMem to be cognizant of Module scope (#876)

Previously, mems marked no dedup would prevent mems with the same
instance name in other modules from deduping

2 files changed, 39 insertions, 19 deletions

diff --git a/src/main/scala/firrtl/passes/memlib/ResolveMaskGranularity.scala b/src/main/scala/firrtl/passes/memlib/ResolveMaskGranularity.scala
index e254dcc9..b552470d 100644
--- a/src/main/scala/firrtl/passes/memlib/ResolveMaskGranularity.scala
+++ b/src/main/scala/firrtl/passes/memlib/ResolveMaskGranularity.scala
@@ -75,12 +75,6 @@ object AnalysisUtils {
        }
     case _ => e
   }
-
-  /** Checks whether the two memories are equivalent in all respects except name
-    */
-  def eqMems(a: DefAnnotatedMemory, b: DefAnnotatedMemory, noDeDupeMems: Seq[String]): Boolean =
-    a == b.copy(info = a.info, name = a.name, memRef = a.memRef) &&
-    !(noDeDupeMems.contains(a.name) || noDeDupeMems.contains(b.name))
 }
 
 /** Determines if a write mask is needed (wmode/en and wmask are equivalent).
diff --git a/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala b/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala
index d195ea55..b0d3731f 100644
--- a/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala
+++ b/src/main/scala/firrtl/passes/memlib/ResolveMemoryReference.scala
@@ -4,7 +4,6 @@ package firrtl.passes
 package memlib
 import firrtl._
 import firrtl.ir._
-import AnalysisUtils.eqMems
 import firrtl.Mappers._
 import firrtl.annotations._
 
@@ -19,30 +18,57 @@ class ResolveMemoryReference extends Transform {
   def inputForm = MidForm
   def outputForm = MidForm
 
-  type AnnotatedMemories = collection.mutable.ArrayBuffer[(String, DefAnnotatedMemory)]
+  /** Helper class for determining when two memories are equivalent while igoring
+    * irrelevant details like name and info
+    */
+  private class WrappedDefAnnoMemory(val underlying: DefAnnotatedMemory) {
+    // Remove irrelevant details for comparison
+    private def generic = underlying.copy(info = NoInfo, name = "", memRef = None)
+    override def hashCode: Int = generic.hashCode
+    override def equals(that: Any): Boolean = that match {
+      case mem: WrappedDefAnnoMemory => this.generic == mem.generic
+      case _ => false
+    }
+  }
+  private def wrap(mem: DefAnnotatedMemory) = new WrappedDefAnnoMemory(mem)
+
+  // Values are Tuple of Module Name and Memory Instance Name
+  private type AnnotatedMemories = collection.mutable.HashMap[WrappedDefAnnoMemory, (String, String)]
+
+  private def dedupable(noDedups: Map[String, Set[String]], module: String, memory: String): Boolean =
+    noDedups.get(module).map(!_.contains(memory)).getOrElse(true)
 
   /** If a candidate memory is identical except for name to another, add an
     *   annotation that references the name of the other memory.
     */
-  def updateMemStmts(mname: String, uniqueMems: AnnotatedMemories, noDeDupeMems: Seq[String])(s: Statement): Statement = s match {
-    case m: DefAnnotatedMemory =>
-      uniqueMems find (x => eqMems(x._2, m, noDeDupeMems)) match {
+  def updateMemStmts(mname: String,
+                     existingMems: AnnotatedMemories,
+                     noDedupMap: Map[String, Set[String]])
+                    (s: Statement): Statement = s match {
+    // If not dedupable, no need to add to existing (since nothing can dedup with it)
+    // We just return the DefAnnotatedMemory as is in the default case below
+    case m: DefAnnotatedMemory if dedupable(noDedupMap, mname, m.name) =>
+      val wrapped = wrap(m)
+      existingMems.get(wrapped) match {
+        case proto @ Some(_) =>
+          m.copy(memRef = proto)
         case None =>
-          uniqueMems += (mname -> m)
+          existingMems(wrapped) = (mname, m.name)
           m
-        case Some((module, proto)) => m copy (memRef = Some(module -> proto.name))
       }
-    case s => s map updateMemStmts(mname, uniqueMems, noDeDupeMems)
+    case s => s.map(updateMemStmts(mname, existingMems, noDedupMap))
   }
 
-  def run(c: Circuit, noDeDupeMems: Seq[String]) = {
-    val uniqueMems = new AnnotatedMemories
-    c copy (modules = c.modules map (m => m map updateMemStmts(m.name, uniqueMems, noDeDupeMems)))
+  def run(c: Circuit, noDedupMap: Map[String, Set[String]]) = {
+    val existingMems = new AnnotatedMemories
+    val modulesx = c.modules.map(m => m.map(updateMemStmts(m.name, existingMems, noDedupMap)))
+    c.copy(modules = modulesx)
   }
   def execute(state: CircuitState): CircuitState = {
     val noDedups = state.annotations.collect {
-      case NoDedupMemAnnotation(ComponentName(cn, _)) => cn
+      case NoDedupMemAnnotation(ComponentName(cn, ModuleName(mn, _))) => mn -> cn
     }
-    state.copy(circuit=run(state.circuit, noDedups))
+    val noDedupMap: Map[String, Set[String]] = noDedups.groupBy(_._1).mapValues(_.map(_._2).toSet)
+    state.copy(circuit = run(state.circuit, noDedupMap))
   }
 }