Cleanup Named Targets (#1311)

* Add CleanupNamedTargets Transform

Adds a new transform, CleanupNamedTargets, that will convert
ReferenceTargets to InstanceTargets if any ReferenceTargets are found
to be actually pointing at instances.

This fixes an issue where a ComponentName pointing at an instance will
be converted to a ReferenceTarget. Later on, this ReferenceTarget will
not be properly renamed by other transforms, e.g.,
EliminateTargetPaths.

Signed-off-by: Schuyler Eldridge <schuyler.eldridge@ibm.com>

* Update ScalaTests

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

4 files changed, 153 insertions, 2 deletions

diff --git a/src/main/scala/firrtl/annotations/transforms/CleanupNamedTargets.scala b/src/main/scala/firrtl/annotations/transforms/CleanupNamedTargets.scala
new file mode 100644
index 00000000..a4219a03
--- /dev/null
+++ b/src/main/scala/firrtl/annotations/transforms/CleanupNamedTargets.scala
@@ -0,0 +1,66 @@
+// See LICENSE for license details.
+
+package firrtl.annotations.transforms
+
+import firrtl._
+import firrtl.annotations.{CircuitTarget, ModuleTarget, MultiTargetAnnotation, ReferenceTarget, SingleTargetAnnotation}
+import firrtl.ir
+import firrtl.options.{Dependency, PreservesAll}
+import firrtl.traversals.Foreachers._
+
+import scala.collection.immutable.{Set => ISet}
+
+/** Replaces all [[firrtl.annotations.ReferenceTarget ReferenceTargets]] pointing at instances with
+  * [[firrtl.annotations.InstanceTarget InstanceTargets]].
+  *
+  * @note This exists because of [[firrtl.annotations.Named Named]] where a [[firrtl.annotations.ComponentName
+  * ComponentName]] is the only way to refer to an instance, but this is resolved incorrectly to a
+  * [[firrtl.annotations.ReferenceTarget ReferenceTarget]].
+  */
+class CleanupNamedTargets extends Transform with DependencyAPIMigration {
+
+  override def prerequisites = Seq(Dependency(passes.RemoveCHIRRTL))
+
+  override def optionalPrerequisites = Seq.empty
+
+  override def optionalPrerequisiteOf = Seq.empty
+
+  override def invalidates(a: Transform) = false
+
+  private def onStatement(statement: ir.Statement)
+                         (implicit references: ISet[ReferenceTarget],
+                          renameMap: RenameMap,
+                          module: ModuleTarget): Unit = statement match {
+    case ir.DefInstance(_, a, b, _) if references(module.instOf(a, b).asReference) =>
+      renameMap.record(module.instOf(a, b).asReference, module.instOf(a, b))
+    case a => statement.foreach(onStatement)
+  }
+
+  private def onModule(module: ir.DefModule)
+                      (implicit references: ISet[ReferenceTarget],
+                       renameMap: RenameMap,
+                       circuit: CircuitTarget): Unit = {
+    implicit val mTarget = circuit.module(module.name)
+    module.foreach(onStatement)
+  }
+
+  override protected def execute(state: CircuitState): CircuitState = {
+
+    implicit val rTargets: ISet[ReferenceTarget] = state.annotations.flatMap {
+      case a: SingleTargetAnnotation[_] => Some(a.target)
+      case a: MultiTargetAnnotation     => a.targets.flatten
+      case _                            => None
+    }.collect {
+      case a: ReferenceTarget => a
+    }.toSet
+
+    implicit val renameMap = RenameMap()
+
+    implicit val cTarget = CircuitTarget(state.circuit.main)
+
+    state.circuit.foreach(onModule)
+
+    state.copy(renames = Some(renameMap))
+  }
+
+}
diff --git a/src/main/scala/firrtl/stage/Forms.scala b/src/main/scala/firrtl/stage/Forms.scala
index 55292fc5..636d0609 100644
--- a/src/main/scala/firrtl/stage/Forms.scala
+++ b/src/main/scala/firrtl/stage/Forms.scala
@@ -20,7 +20,8 @@ object Forms {
     Seq( Dependency(passes.CheckChirrtl),
          Dependency(passes.CInferTypes),
          Dependency(passes.CInferMDir),
-         Dependency(passes.RemoveCHIRRTL) )
+         Dependency(passes.RemoveCHIRRTL),
+         Dependency[annotations.transforms.CleanupNamedTargets] )
 
   val WorkingIR: Seq[TransformDependency] = MinimalHighForm :+ Dependency(passes.ToWorkingIR)
 
diff --git a/src/test/scala/firrtlTests/LoweringCompilersSpec.scala b/src/test/scala/firrtlTests/LoweringCompilersSpec.scala
index 802596c5..e1de7e34 100644
--- a/src/test/scala/firrtlTests/LoweringCompilersSpec.scala
+++ b/src/test/scala/firrtlTests/LoweringCompilersSpec.scala
@@ -132,7 +132,10 @@ class LoweringCompilersSpec extends AnyFlatSpec with Matchers {
 
   it should "replicate the old order" in {
     val tm = new TransformManager(Forms.MinimalHighForm, Forms.ChirrtlForm)
-    compare(legacyTransforms(new firrtl.ChirrtlToHighFirrtl), tm)
+    val patches = Seq(
+      Add(5, Seq(Dependency[firrtl.annotations.transforms.CleanupNamedTargets]))
+    )
+    compare(legacyTransforms(new firrtl.ChirrtlToHighFirrtl), tm, patches)
   }
 
   behavior of "IRToWorkingIR"
diff --git a/src/test/scala/firrtlTests/annotationTests/CleanupNamedTargetsSpec.scala b/src/test/scala/firrtlTests/annotationTests/CleanupNamedTargetsSpec.scala
new file mode 100644
index 00000000..58cb3d11
--- /dev/null
+++ b/src/test/scala/firrtlTests/annotationTests/CleanupNamedTargetsSpec.scala
@@ -0,0 +1,81 @@
+// See LICENSE for license details.
+
+package firrtlTests.annotationTests
+
+import firrtl.{AnnotationSeq, CircuitState, Parser, UnknownForm}
+import firrtl.annotations.{
+  CircuitName,
+  CircuitTarget,
+  ComponentName,
+  ModuleName,
+  MultiTargetAnnotation,
+  ReferenceTarget,
+  SingleTargetAnnotation,
+  Target}
+import firrtl.annotations.transforms.CleanupNamedTargets
+
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+object CleanupNamedTargetsSpec {
+
+  case class SingleReferenceAnnotation(target: ReferenceTarget) extends SingleTargetAnnotation[ReferenceTarget] {
+    override def duplicate(a: ReferenceTarget) = this.copy(target = a)
+  }
+
+  case class MultiReferenceAnnotation(targets: Seq[Seq[Target]]) extends MultiTargetAnnotation {
+    override def duplicate(a: Seq[Seq[Target]]) = this.copy(targets = a)
+  }
+
+}
+
+class CleanupNamedTargetsSpec extends AnyFlatSpec with Matchers {
+
+  import CleanupNamedTargetsSpec._
+
+  class F {
+
+    val transform = new CleanupNamedTargets
+
+    val circuit =
+      """|circuit Foo:
+         |  module Bar:
+         |    skip
+         |  module Foo:
+         |    inst bar of Bar
+         |    wire baz: UInt<1>
+         |""".stripMargin
+
+    lazy val foo = CircuitTarget("Foo").module("Foo")
+
+    lazy val barTarget = ComponentName("bar", ModuleName("Foo", CircuitName("Foo"))).toTarget
+
+    lazy val bazTarget = ComponentName("baz", ModuleName("Foo", CircuitName("Foo"))).toTarget
+
+    def circuitState(a: AnnotationSeq): CircuitState = CircuitState(Parser.parse(circuit), UnknownForm, a, None)
+
+  }
+
+  behavior of "CleanupNamedTargets"
+
+  it should "convert a SingleTargetAnnotation[ReferenceTarget] of an instance to an InstanceTarget" in new F {
+    val annotations: AnnotationSeq = Seq(SingleReferenceAnnotation(barTarget))
+
+    transform.transform(circuitState(annotations)).renames.get.get(barTarget) should be {
+      Some(Seq(foo.instOf("bar", "Bar")))
+    }
+  }
+
+  it should "convert a MultiTargetAnnotation with a ReferenceTarget of an instance to an InstanceTarget" in new F {
+    val annotations: AnnotationSeq = Seq(MultiReferenceAnnotation(Seq(Seq(barTarget), Seq(bazTarget))))
+
+    val renames = transform.transform(circuitState(annotations)).renames.get
+
+    renames.get(barTarget) should be (Some(Seq(foo.instOf("bar", "Bar"))))
+
+    info("and not touch a true ReferenceAnnotation")
+    renames.get(bazTarget) should be (None)
+
+  }
+
+}