Refactor RenameMap to rename Components if their Module is renamed

2 files changed, 130 insertions, 15 deletions

diff --git a/src/main/scala/firrtl/Compiler.scala b/src/main/scala/firrtl/Compiler.scala
index b49af505..43596765 100644
--- a/src/main/scala/firrtl/Compiler.scala
+++ b/src/main/scala/firrtl/Compiler.scala
@@ -10,10 +10,7 @@ import scala.collection.mutable
 import firrtl.annotations._  // Note that wildcard imports are not great....
 import firrtl.ir.Circuit
 import firrtl.Utils.{error, throwInternalError}
-/**
- * RenameMap maps old names to modified names.  Generated by transformations
- * that modify names
- */
+
 object RenameMap {
   def apply(map: Map[Named, Seq[Named]]) = {
     val rm = new RenameMap
@@ -22,8 +19,52 @@ object RenameMap {
   }
   def apply() = new RenameMap
 }
-class RenameMap {
-  val renameMap = new mutable.HashMap[Named, Seq[Named]]()
+/** Map old names to new names
+  *
+  * Transforms that modify names should return a [[RenameMap]] with the [[CircuitState]]
+  * These are mutable datastructures for convenience
+  */
+// TODO This should probably be refactored into immutable and mutable versions
+final class RenameMap private () {
+  private val underlying = mutable.HashMap[Named, Seq[Named]]()
+  /** Get new names for an old name
+    *
+    * This is analogous to get on standard Scala collection Maps
+    * None indicates the key was not renamed
+    * Empty indicates the name was deleted
+    */
+  // TODO Is there a better way to express this?
+  def get(key: Named): Option[Seq[Named]] = {
+    underlying.get(key) match {
+      // If the key was renamed, check if anything it renamed to is a component
+      // If so, check if nested modules were renamed
+      case Some(names) => Some(names.flatMap {
+        case comp @ ComponentName(cname, mod) =>
+          underlying.get(mod) match {
+            case Some(mods) => mods.map {
+              case modx: ModuleName =>
+                ComponentName(cname, modx)
+              case _ => error("Unexpected rename of Module to non-Module!")
+            }
+            case None => List(comp)
+          }
+        case other => List(other)
+      })
+      // If key wans't renamed, still check if it's a component
+      // If so, check if nexted modules were renamed
+      case None => key match {
+        case ComponentName(cname, mod) =>
+          underlying.get(mod).map(_.map {
+            case modx: ModuleName =>
+              ComponentName(cname, modx)
+            case _ => error("Unexpected rename of Module to non-Module!")
+          })
+        case other => None
+      }
+    }
+  }
+
+  // Mutable helpers
   private var circuitName: String = ""
   private var moduleName: String = ""
   def setModule(s: String) =
@@ -40,18 +81,18 @@ class RenameMap {
   }
   def rename(from: Named, to: Named): Unit = rename(from, Seq(to))
   def rename(from: Named, tos: Seq[Named]): Unit = (from, tos) match {
-    case (x, Seq(y)) if x == y =>
+    case (x, Seq(y)) if x == y => // TODO is this check expensive in common case?
     case _ =>
-      renameMap(from) = renameMap.getOrElse(from, Seq.empty) ++ tos
+      underlying(from) = underlying.getOrElse(from, Seq.empty) ++ tos
   }
   def delete(names: Seq[String]): Unit = names.foreach(delete(_))
   def delete(name: String): Unit =
     delete(ComponentName(name, ModuleName(moduleName, CircuitName(circuitName))))
-  def delete(name: Named): Unit = 
-    renameMap(name) = Seq.empty
+  def delete(name: Named): Unit =
+    underlying(name) = Seq.empty
   def addMap(map: Map[Named, Seq[Named]]) =
-    renameMap ++= map
-  def serialize: String = renameMap.map { case (k, v) =>
+    underlying ++= map
+  def serialize: String = underlying.map { case (k, v) =>
     k.serialize + "=>" + v.map(_.serialize).mkString(", ")
   }.mkString("\n")
 }
@@ -151,7 +192,7 @@ final case object MidForm extends CircuitForm(1)
   */
 final case object LowForm extends CircuitForm(0)
 /** Unknown Form
-  * 
+  *
   * Often passes may modify a circuit (e.g. InferTypes), but return
   * a circuit in the same form it was given.
   *
@@ -239,10 +280,10 @@ abstract class Transform extends LazyLogging {
     }
 
     // For each annotation, rename all annotations.
-    val renames = renameOpt.getOrElse(RenameMap()).renameMap
+    val renames = renameOpt.getOrElse(RenameMap())
     for {
       anno <- newAnnotations.toSeq
-      newAnno <- anno.update(renames.getOrElse(anno.target, Seq(anno.target)))
+      newAnno <- anno.update(renames.get(anno.target).getOrElse(Seq(anno.target)))
     } yield newAnno
   }
 }
diff --git a/src/test/scala/firrtlTests/RenameMapSpec.scala b/src/test/scala/firrtlTests/RenameMapSpec.scala
new file mode 100644
index 00000000..9d19bb72
--- /dev/null
+++ b/src/test/scala/firrtlTests/RenameMapSpec.scala
@@ -0,0 +1,74 @@
+// See LICENSE for license details.
+
+package firrtlTests
+
+import firrtl.RenameMap
+import firrtl.annotations.{
+  CircuitName,
+  ModuleName,
+  ComponentName
+}
+
+class RenameMapSpec extends FirrtlFlatSpec {
+  val cir = CircuitName("Top")
+  val modA = ModuleName("A", cir)
+  val modB = ModuleName("B", cir)
+  val foo = ComponentName("foo", modA)
+  val bar = ComponentName("bar", modA)
+  val fizz = ComponentName("fizz", modA)
+  val fooB = ComponentName("foo", modB)
+  val barB = ComponentName("bar", modB)
+
+  behavior of "RenameMap"
+
+  it should "return None if it does not rename something" in {
+    val renames = RenameMap()
+    renames.get(modA) should be (None)
+    renames.get(foo) should be (None)
+  }
+
+  it should "return a Seq of renamed things if it does rename something" in {
+    val renames = RenameMap()
+    renames.rename(foo, bar)
+    renames.get(foo) should be (Some(Seq(bar)))
+  }
+
+  it should "allow something to be renamed to multiple things" in {
+    val renames = RenameMap()
+    renames.rename(foo, bar)
+    renames.rename(foo, fizz)
+    renames.get(foo) should be (Some(Seq(bar, fizz)))
+  }
+
+  it should "allow something to be renamed to nothing (ie. deleted)" in {
+    val renames = RenameMap()
+    renames.rename(foo, Seq())
+    renames.get(foo) should be (Some(Seq()))
+  }
+
+  it should "return None if something is renamed to itself" in {
+    val renames = RenameMap()
+    renames.rename(foo, foo)
+    renames.get(foo) should be (None)
+  }
+
+  it should "allow components to change module" in {
+    val renames = RenameMap()
+    renames.rename(foo, fooB)
+    renames.get(foo) should be (Some(Seq(fooB)))
+  }
+
+  it should "rename components if their module is renamed" in {
+    val renames = RenameMap()
+    renames.rename(modA, modB)
+    renames.get(foo) should be (Some(Seq(fooB)))
+    renames.get(bar) should be (Some(Seq(barB)))
+  }
+
+  it should "rename renamed components if the module of the target component is renamed" in {
+    val renames = RenameMap()
+    renames.rename(modA, modB)
+    renames.rename(foo, bar)
+    renames.get(foo) should be (Some(Seq(barB)))
+  }
+}