Implement InstanceTarget Behavior for Dedup + EliminateTargetPaths (#1539)

- RenameMap Behavior
-- Prevent transitive renaming A -> B -> C (continueRenaming)
-- Prevent transitive renaming for self-renames

- Target
-- Override toString as serialize for CompleteTarget
-- Expansion of stripHierarchy to enable stripping InstanceTargets to become ModuleTargets
Annotations
-- Bugfix in extractComponents where Products were not iterated over
-- Converts renamed targets to local targets using Target.referringModule to preserve sticky behavior

- Eliminate Target Paths
-- Make DuplicationHelper use LinkedHashMap, as we iterate over its contents and convert to Seq in def makePathless
-- Add DupedResult to map original module to new module targets
-- Update renaming to record a map from all relative instance paths to original module, to new module target
-- Consumes DedupedResult to give better name to new duplicated module if it was originally deduplicated
-- Reorder modules in attempt to preserve original ordering, pre-deduplication
-- Move utility functions to object
-- Bugfix: add self-renames to prevent ofModule _ of target _ cannot be renamed to Vector(_, _, _, ...) errors

- Dedup
-- Changed NoDedupAnnotation to contain ModuleTarget, rather than ModuleName
-- Added DedupedResult to map original module to the duplicate module
-- Consumes DupedResult to pick better name, if it existed
-- Updates renaming to chain the following: instancify deduped modules, remap differently named internal signals, then remap AST modules
-- Move utility functions to object
-- Remove annotations as part of determination of dedup correctness
-- Bugfix: add instance renames so that deduped modules have their instances properly renamed

- Dead Code Elimination
-- Add deletion of ASTModules

- Tests
-- Morphism Spec to ensure Dedup -> EliminateTargetPaths and EliminateTargetPaths -> Dedup patterns work properly
-- Update existing tests to make sure they work properly
-- Add Dedup tests to demonstrate instance renaming bug, EliminateTargetPaths for ofModule rename bug, and update RenameMap tests

Co-authored-by: Schuyler Eldridge <schuyler.eldridge@ibm.com>
Co-authored-by: Adam Izraelevitz <adam.izraelevitz@sifive.com>
Co-authored-by: Adam Izraelevitz <azidar@gmail.com>
Co-authored-by: Jack Koenig <koenig@sifive.com>

4 files changed, 206 insertions, 55 deletions

diff --git a/src/main/scala/firrtl/annotations/Annotation.scala b/src/main/scala/firrtl/annotations/Annotation.scala
index 4c39bfee..fcbcdc96 100644
--- a/src/main/scala/firrtl/annotations/Annotation.scala
+++ b/src/main/scala/firrtl/annotations/Annotation.scala
@@ -27,7 +27,8 @@ trait Annotation extends Product {
   private def extractComponents(ls: scala.collection.Traversable[_]): Seq[Target] = {
     ls.collect {
       case c: Target => Seq(c)
-      case ls: scala.collection.Traversable[_] => extractComponents(ls)
+      case o: Product => extractComponents(o.productIterator.toIterable)
+      case x: scala.collection.Traversable[_] => extractComponents(x)
     }.foldRight(Seq.empty[Target])((seq, c) => c ++ seq)
   }
 
@@ -59,29 +60,39 @@ trait SingleTargetAnnotation[T <: Named] extends Annotation {
       case c: Target =>
         val x = renames.get(c)
         x.map(newTargets => newTargets.map(t => duplicate(t.asInstanceOf[T]))).getOrElse(List(this))
-      case _: Named =>
+      case from: Named =>
         val ret = renames.get(Target.convertNamed2Target(target))
-        ret.map(_.map(newT => Target.convertTarget2Named(newT: @unchecked) match {
-          case newTarget: T @unchecked =>
-            try {
-              duplicate(newTarget)
-            }
-            catch {
-              case _: java.lang.ClassCastException =>
-                val msg = s"${this.getClass.getName} target ${target.getClass.getName} " +
-                  s"cannot be renamed to ${newTarget.getClass}"
-                throw AnnotationException(msg)
-            }
-        })).getOrElse(List(this))
+        ret.map(_.map { newT =>
+          val result = newT match {
+            case c: InstanceTarget => ModuleName(c.ofModule, CircuitName(c.circuit))
+            case c: IsMember =>
+              val local = Target.referringModule(c)
+              c.setPathTarget(local)
+            case c: CircuitTarget => c.toNamed
+            case other => throw Target.NamedException(s"Cannot convert $other to [[Named]]")
+          }
+          Target.convertTarget2Named(result) match {
+            case newTarget: T @unchecked =>
+              try {
+                duplicate(newTarget)
+              }
+              catch {
+                case _: java.lang.ClassCastException =>
+                  val msg = s"${this.getClass.getName} target ${target.getClass.getName} " +
+                    s"cannot be renamed to ${newTarget.getClass}"
+                  throw AnnotationException(msg)
+              }
+          }
+        }).getOrElse(List(this))
     }
   }
 }
 
 /** [[MultiTargetAnnotation]] keeps the renamed targets grouped within a single annotation. */
 trait MultiTargetAnnotation extends Annotation {
-  /** Contains a sequence of [[Target]].
+  /** Contains a sequence of targets.
     * When created, [[targets]] should be assigned by `Seq(Seq(TargetA), Seq(TargetB), Seq(TargetC))`
-    * */
+    */
   val targets: Seq[Seq[Target]]
 
   /** Create another instance of this Annotation*/
diff --git a/src/main/scala/firrtl/annotations/Target.scala b/src/main/scala/firrtl/annotations/Target.scala
index 10c74e77..f33a8fdf 100644
--- a/src/main/scala/firrtl/annotations/Target.scala
+++ b/src/main/scala/firrtl/annotations/Target.scala
@@ -366,6 +366,9 @@ trait CompleteTarget extends Target {
   def addHierarchy(root: String, instance: String): IsComponent
 
   override def toTarget: CompleteTarget = this
+
+  // Very useful for debugging, I (@azidar) think this is reasonable
+  override def toString = serialize
 }
 
 
@@ -668,10 +671,14 @@ case class InstanceTarget(circuit: String,
   override def instOf(inst: String, of: String): InstanceTarget = InstanceTarget(circuit, module, asPath, inst, of)
 
   override def stripHierarchy(n: Int): IsModule = {
-    require(path.size >= n, s"Cannot strip $n levels of hierarchy from $this")
+    require(path.size + 1 >= n, s"Cannot strip $n levels of hierarchy from $this")
     if(n == 0) this else {
-      val newModule = path(n - 1)._2.value
-      InstanceTarget(circuit, newModule, path.drop(n), instance, ofModule)
+      if(path.size < n){
+        ModuleTarget(circuit, ofModule)
+      } else {
+        val newModule = path(n - 1)._2.value
+        InstanceTarget(circuit, newModule, path.drop(n), instance, ofModule)
+      }
     }
   }
 
diff --git a/src/main/scala/firrtl/annotations/analysis/DuplicationHelper.scala b/src/main/scala/firrtl/annotations/analysis/DuplicationHelper.scala
index f892c508..8f925ee7 100644
--- a/src/main/scala/firrtl/annotations/analysis/DuplicationHelper.scala
+++ b/src/main/scala/firrtl/annotations/analysis/DuplicationHelper.scala
@@ -12,24 +12,25 @@ import scala.collection.mutable
   * Calculates needed modifications to a circuit's module/instance hierarchy
   */
 case class DuplicationHelper(existingModules: Set[String]) {
+
   // Maps instances to the module it instantiates (an ofModule)
-  type InstanceOfModuleMap = mutable.HashMap[Instance, OfModule]
+  type InstanceOfModuleMap = mutable.LinkedHashMap[Instance, OfModule]
 
   // Maps a module to the instance/ofModules it instantiates
-  type ModuleHasInstanceOfModuleMap = mutable.HashMap[String, InstanceOfModuleMap]
+  type ModuleHasInstanceOfModuleMap = mutable.LinkedHashMap[String, InstanceOfModuleMap]
 
   // Maps original module names to new duplicated modules and their encapsulated instance/ofModules
-  type DupMap = mutable.HashMap[String, ModuleHasInstanceOfModuleMap]
+  type DupMap = mutable.LinkedHashMap[String, ModuleHasInstanceOfModuleMap]
 
   // Internal state to keep track of how paths duplicate
   private val dupMap = new DupMap()
 
   // Internal record of which paths are renamed to which new names, in the case of a collision
-  private val cachedNames = mutable.HashMap[(String, Seq[(Instance, OfModule)]), String]() ++
+  private val cachedNames = mutable.LinkedHashMap[(String, Seq[(Instance, OfModule)]), String]() ++
     existingModules.map(m => (m, Nil) -> m)
 
   // Internal record of all paths to ensure unique name generation
-  private val allModules = mutable.HashSet[String]() ++ existingModules
+  private val allModules = mutable.LinkedHashSet[String]() ++ existingModules
 
   /** Updates internal state (dupMap) to calculate instance hierarchy modifications so t's tokens in an instance can be
     * expressed as a tokens in a module (e.g. uniquify/duplicate the instance path in t's tokens)
diff --git a/src/main/scala/firrtl/annotations/transforms/EliminateTargetPaths.scala b/src/main/scala/firrtl/annotations/transforms/EliminateTargetPaths.scala
index a4cd2f3d..6bafa071 100644
--- a/src/main/scala/firrtl/annotations/transforms/EliminateTargetPaths.scala
+++ b/src/main/scala/firrtl/annotations/transforms/EliminateTargetPaths.scala
@@ -9,9 +9,10 @@ import firrtl.annotations.TargetToken.{Instance, OfModule, fromDefModuleToTarget
 import firrtl.annotations.analysis.DuplicationHelper
 import firrtl.annotations._
 import firrtl.ir._
-import firrtl.{CircuitState, DependencyAPIMigration, FirrtlInternalException, RenameMap, Transform, WDefInstance}
+import firrtl.{AnnotationSeq, CircuitState, DependencyAPIMigration, FirrtlInternalException, RenameMap, Transform}
 import firrtl.options.PreservesAll
 import firrtl.stage.Forms
+import firrtl.transforms.DedupedResult
 
 import scala.collection.mutable
 
@@ -26,8 +27,66 @@ case class ResolvePaths(targets: Seq[CompleteTarget]) extends Annotation {
   }
 }
 
+/** Holds the mapping from original module to the new, duplicated modules
+  * The original module target is unaffected by renaming
+  * @param newModules Instance target of what the original module now points to
+  * @param originalModule Original module
+  */
+case class DupedResult(newModules: Set[IsModule], originalModule: ModuleTarget) extends MultiTargetAnnotation {
+  override val targets: Seq[Seq[Target]] = Seq(newModules.toSeq)
+  override def duplicate(n: Seq[Seq[Target]]): Annotation = {
+    n.toList match {
+      case Seq(newMods) => DupedResult(newMods.collect { case x: IsModule => x }.toSet, originalModule)
+      case _ => DupedResult(Set.empty, originalModule)
+    }
+  }
+}
+
 case class NoSuchTargetException(message: String) extends FirrtlInternalException(message)
 
+object EliminateTargetPaths {
+
+  def renameModules(c: Circuit, toRename: Map[String, String], renameMap: RenameMap): Circuit = {
+    val ct = CircuitTarget(c.main)
+    val cx = if(toRename.contains(c.main)) {
+      renameMap.record(ct, CircuitTarget(toRename(c.main)))
+      c.copy(main = toRename(c.main))
+    } else {
+      c
+    }
+    def onMod(m: DefModule): DefModule = {
+      m map onStmt match {
+        case e: ExtModule if toRename.contains(e.name) =>
+          renameMap.record(ct.module(e.name), ct.module(toRename(e.name)))
+          e.copy(name = toRename(e.name))
+        case e: Module if toRename.contains(e.name)    =>
+          renameMap.record(ct.module(e.name), ct.module(toRename(e.name)))
+          e.copy(name = toRename(e.name))
+        case o => o
+      }
+    }
+    def onStmt(s: Statement): Statement = s map onStmt match {
+      case w@DefInstance(info, name, module, _) if toRename.contains(module) => w.copy(module = toRename(module))
+      case other => other
+    }
+    cx map onMod
+  }
+
+  def reorderModules(c: Circuit, toReorder: Map[String, Double]): Circuit = {
+    val newOrderMap = c.modules.zipWithIndex.map {
+      case (m, _) if toReorder.contains(m.name) => m.name -> toReorder(m.name)
+      case (m, i) if c.modules.size > 1 => m.name -> i.toDouble / (c.modules.size - 1)
+      case (m, _) => m.name -> 1.0
+    }.toMap
+
+    val newOrder = c.modules.sortBy { m => newOrderMap(m.name) }
+
+    c.copy(modules = newOrder)
+  }
+
+
+}
+
 /** For a set of non-local targets, modify the instance/module hierarchy of the circuit such that
   * the paths in each non-local target can be removed
   *
@@ -44,6 +103,7 @@ case class NoSuchTargetException(message: String) extends FirrtlInternalExceptio
   * C/x -> (C/x, C_/x) // where x is any reference in C
   */
 class EliminateTargetPaths extends Transform with DependencyAPIMigration with PreservesAll[Transform] {
+  import EliminateTargetPaths._
 
   override def prerequisites = Forms.MinimalHighForm
   override def optionalPrerequisites = Seq.empty
@@ -62,9 +122,6 @@ class EliminateTargetPaths extends Transform with DependencyAPIMigration with Pr
     case d@DefInstance(_, name, module, _) =>
       val ofModule = dupMap.getNewOfModule(originalModule, newModule, Instance(name), OfModule(module)).value
       d.copy(module = ofModule)
-    case d@WDefInstance(_, name, module, _) =>
-      val ofModule = dupMap.getNewOfModule(originalModule, newModule, Instance(name), OfModule(module)).value
-      d.copy(module = ofModule)
     case other => other map onStmt(dupMap)(originalModule, newModule)
   }
 
@@ -73,7 +130,10 @@ class EliminateTargetPaths extends Transform with DependencyAPIMigration with Pr
     * @param targets
     * @return
     */
-  def run(cir: Circuit, targets: Seq[IsMember]): (Circuit, RenameMap) = {
+  def run(cir: Circuit,
+          targets: Seq[IsMember],
+          iGraph: InstanceGraph
+         ): (Circuit, RenameMap, AnnotationSeq) = {
 
     val dupMap = DuplicationHelper(cir.modules.map(_.name).toSet)
 
@@ -85,8 +145,11 @@ class EliminateTargetPaths extends Transform with DependencyAPIMigration with Pr
 
     // Foreach module, calculate the unique names of its duplicates
     // Then, update the ofModules of instances that it encapsulates
-    cir.modules.foreach { m =>
-      dupMap.getDuplicates(m.name).foreach { newName =>
+
+    val ct = CircuitTarget(cir.main)
+    val annos = cir.modules.map { m =>
+      val newNames = dupMap.getDuplicates(m.name)
+      newNames.foreach { newName =>
         val newM = m match {
           case e: ExtModule => e.copy(name = newName)
           case o: Module =>
@@ -94,6 +157,7 @@ class EliminateTargetPaths extends Transform with DependencyAPIMigration with Pr
         }
         duplicatedModuleList += newM
       }
+      DupedResult(newNames.map(ct.module), ct.module(m.name))
     }
 
     val finalModuleList = duplicatedModuleList
@@ -105,34 +169,74 @@ class EliminateTargetPaths extends Transform with DependencyAPIMigration with Pr
     /* Foreach target, calculate the pathless version and only rename targets that are instantiated. Additionally, rename
      * module targets
      */
+    def addRecord(old: IsMember, newPathless: IsMember): Unit = old match {
+      case x: ModuleTarget =>
+        renameMap.record(x, newPathless)
+      case x: IsComponent if x.path.isEmpty =>
+        renameMap.record(x, newPathless)
+      case x: IsComponent =>
+        renameMap.record(x, newPathless)
+        addRecord(x.stripHierarchy(1), newPathless)
+    }
+    val duplicatedParents = mutable.Set[OfModule]()
     targets.foreach { t =>
-      val newTsx = dupMap.makePathless(t)
-      val newTs = newTsx
-      if(newTs.nonEmpty) {
-        renameMap.record(t, newTs)
-        val m = Target.referringModule(t)
-        val duplicatedModules = newTs.map(Target.referringModule)
-        val oldModule: Option[ModuleTarget] = m match {
-          case a: ModuleTarget if finalModuleSet(a.module) => Some(a)
-          case _                                           => None
+      val newTs = dupMap.makePathless(t)
+      val path = t.asPath
+      if (path.nonEmpty) { duplicatedParents += path(0)._2 }
+      newTs.toList match {
+        case Seq(pathless) =>
+          val mt = Target.referringModule(pathless)
+          addRecord(t, pathless)
+          renameMap.record(Target.referringModule(t), mt)
+        case _ =>
+      }
+    }
+
+    def addSelfRecord(mod: IsModule): Unit = mod match {
+      case m: ModuleTarget =>
+      case i: InstanceTarget if renameMap.underlying.contains(i) =>
+      case i: InstanceTarget =>
+        renameMap.record(i, i)
+        addSelfRecord(i.stripHierarchy(1))
+    }
+    val topMod = ModuleTarget(cir.main, cir.main)
+    duplicatedParents.foreach { parent =>
+      val paths = iGraph.findInstancesInHierarchy(parent.value)
+      val newTargets = paths.map { path =>
+        path.tail.foldLeft(topMod: IsModule) { case (mod, wDefInst) =>
+          mod.instOf(wDefInst.name, wDefInst.module)
         }
-        renameMap.record(m, (duplicatedModules).distinct)
       }
+      newTargets.foreach(addSelfRecord(_))
     }
 
     // Return modified circuit and associated renameMap
-    (cir.copy(modules = finalModuleList), renameMap)
+    (cir.copy(modules = finalModuleList), renameMap, annos)
   }
 
   override def execute(state: CircuitState): CircuitState = {
+    val moduleNames = state.circuit.modules.map(_.name).toSet
+
+    val (remainingAnnotations, targetsToEliminate, previouslyDeduped) =
+      state.annotations.foldLeft(
+        ( Vector.empty[Annotation],
+          Seq.empty[CompleteTarget],
+          Map.empty[IsModule, (ModuleTarget, Double)]
+        )
+      ) { case ((remainingAnnos, targets, dedupedResult), anno)  =>
+          anno match {
+            case ResolvePaths(ts)          =>
+              (remainingAnnos, ts ++ targets, dedupedResult)
+            case DedupedResult(orig, dups, idx) if dups.nonEmpty =>
+              (remainingAnnos, targets, dedupedResult ++ dups.map(_ -> (orig, idx)).toMap)
+            case other =>
+              (remainingAnnos :+ other, targets, dedupedResult)
+          }
+      }
 
-    val (annotations, annotationsx) = state.annotations.partition{
-      case a: ResolvePaths => true
-      case _               => false
-    }
 
     // Collect targets that are not local
-    val targets = annotations.map(_.asInstanceOf[ResolvePaths]).flatMap(_.targets.collect { case x: IsMember => x })
+    val targets = targetsToEliminate.collect { case x: IsMember => x }
 
     // Check validity of paths in targets
     val iGraph = new InstanceGraph(state.circuit)
@@ -140,7 +244,7 @@ class EliminateTargetPaths extends Transform with DependencyAPIMigration with Pr
     val targetsWithInvalidPaths = mutable.ArrayBuffer[IsMember]()
     targets.foreach { t =>
       val path = t match {
-        case m: ModuleTarget => Nil
+        case _: ModuleTarget => Nil
         case i: InstanceTarget => i.asPath
         case r: ReferenceTarget => r.path
       }
@@ -157,14 +261,18 @@ class EliminateTargetPaths extends Transform with DependencyAPIMigration with Pr
       throw NoSuchTargetException(s"""Some targets have illegal paths that cannot be resolved/eliminated: $string""")
     }
 
-    // get rid of path prefixes of modules with only one instance so we don't rename them
+    /* get rid of path prefixes of modules with only one instance so we don't rename them
+     * If instance targeted is in fact the only instance of a module, then it should not be renamed
+     * E.g. if Eliminate Target Paths on ~Top|Top/foo:Foo, but that is the only instance of Foo, then should return
+     *   ~Top|Top/foo:Foo, not ~Top|Top/foo:Foo___Top_foo
+     */
     val isSingleInstMod: String => Boolean = {
       val cache = mutable.Map.empty[String, Boolean]
       mod => cache.getOrElseUpdate(mod, iGraph.findInstancesInHierarchy(mod).size == 1)
     }
     val firstRenameMap = RenameMap()
-    val nonSingletonTargets = targets.foldLeft(Seq.empty[IsMember]) {
-      case (acc, t: IsComponent) if t.asPath.nonEmpty =>
+    val nonSingletonTargets = targets.foldRight(Seq.empty[IsMember]) {
+      case (t: IsComponent, acc) if t.asPath.nonEmpty =>
         val origPath = t.asPath
         val (singletonPrefix, rest) = origPath.partition {
           case (_, OfModule(mod)) =>
@@ -191,14 +299,14 @@ class EliminateTargetPaths extends Transform with DependencyAPIMigration with Pr
         } else {
           t +: acc
         }
-      case (acc, t) => t +: acc
+      case (t, acc) => t +: acc
     }
 
-    val (newCircuit, nextRenameMap) = run(state.circuit, nonSingletonTargets)
+    val (newCircuit, nextRenameMap, newAnnos) = run(state.circuit, nonSingletonTargets, iGraph)
 
     val renameMap =
       if (firstRenameMap.hasChanges) {
-        firstRenameMap andThen nextRenameMap
+        firstRenameMap.andThen(nextRenameMap)
       } else {
         nextRenameMap
       }
@@ -217,6 +325,30 @@ class EliminateTargetPaths extends Transform with DependencyAPIMigration with Pr
       newCircuit.copy(modules = modulesx)
     }
 
-    state.copy(circuit = newCircuitGC, renames = Some(renameMap), annotations = annotationsx)
+    val renamedModuleMap = RenameMap()
+
+    // If previous instance target mapped to a single previously deduped module, return original name
+    // E.g. if previously ~Top|Top/foo:Foo was deduped to ~Top|Top/foo:Bar, then
+    //  Eliminate target paths on ~Top|Top/foo:Bar should rename to ~Top|Top/foo:Foo, not
+    //  ~Top|Top/foo:Bar___Top_foo
+    val newModuleNameMapping = previouslyDeduped.flatMap {
+      case (current: IsModule, (orig: ModuleTarget, idx)) =>
+        renameMap.get(current).collect { case Seq(ModuleTarget(_, m)) => m -> orig.name }
+    }
+
+    val renamedCircuit = renameModules(newCircuitGC, newModuleNameMapping, renamedModuleMap)
+
+    val reorderedCircuit = reorderModules(renamedCircuit,
+      previouslyDeduped.map {
+        case (current: IsModule, (orig: ModuleTarget, idx)) =>
+          orig.name -> idx
+      }
+    )
+
+    state.copy(
+      circuit = reorderedCircuit,
+      renames = Some(renameMap.andThen(renamedModuleMap)),
+      annotations = remainingAnnotations ++ newAnnos
+    )
   }
 }