Provide an annotation mix-in that marks RTs as dontTouch (#1433)

* Provide an annotation mix-in that marks RTs as dontTouch

* Update src/main/scala/firrtl/transforms/OptimizationAnnotations.scala

Co-Authored-By: Albert Magyar <albert.magyar@gmail.com>

* Update src/test/scala/firrtlTests/DCETests.scala

Co-Authored-By: Albert Magyar <albert.magyar@gmail.com>

* Update src/main/scala/firrtl/transforms/OptimizationAnnotations.scala

* Update OptimizationAnnotations.scala

Co-authored-by: Albert Magyar <albert.magyar@gmail.com>

5 files changed, 71 insertions, 6 deletions

diff --git a/src/main/scala/firrtl/transforms/ConstantPropagation.scala b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
index 201c3325..55c897b3 100644
--- a/src/main/scala/firrtl/transforms/ConstantPropagation.scala
+++ b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
@@ -690,8 +690,12 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
   }
 
   def execute(state: CircuitState): CircuitState = {
-    val dontTouches: Seq[(OfModule, String)] = state.annotations.collect {
-      case DontTouchAnnotation(Target(_, Some(m), Seq(Ref(c)))) => m.OfModule -> c
+    val dontTouchRTs = state.annotations.flatMap {
+      case anno: HasDontTouches => anno.dontTouches
+      case o => Nil
+    }
+    val dontTouches: Seq[(OfModule, String)] = dontTouchRTs.map {
+      case Target(_, Some(m), Seq(Ref(c))) => m.OfModule -> c
     }
     // Map from module name to component names
     val dontTouchMap: Map[OfModule, Set[String]] =
diff --git a/src/main/scala/firrtl/transforms/DeadCodeElimination.scala b/src/main/scala/firrtl/transforms/DeadCodeElimination.scala
index 308d68df..983f1048 100644
--- a/src/main/scala/firrtl/transforms/DeadCodeElimination.scala
+++ b/src/main/scala/firrtl/transforms/DeadCodeElimination.scala
@@ -337,8 +337,9 @@ class DeadCodeElimination extends Transform with ResolvedAnnotationPaths with Re
     Seq(classOf[DontTouchAnnotation], classOf[OptimizableExtModuleAnnotation])
 
   def execute(state: CircuitState): CircuitState = {
-    val dontTouches: Seq[LogicNode] = state.annotations.collect {
-      case DontTouchAnnotation(component: ReferenceTarget) if component.isLocal => LogicNode(component)
+    val dontTouches: Seq[LogicNode] = state.annotations.flatMap {
+      case anno: HasDontTouches => anno.dontTouches.filter(_.isLocal).map(LogicNode(_))
+      case o => Nil
     }
     val doTouchExtMods: Seq[String] = state.annotations.collect {
       case OptimizableExtModuleAnnotation(ModuleName(name, _)) => name
diff --git a/src/main/scala/firrtl/transforms/OptimizationAnnotations.scala b/src/main/scala/firrtl/transforms/OptimizationAnnotations.scala
index b0cf2f02..ff44afec 100644
--- a/src/main/scala/firrtl/transforms/OptimizationAnnotations.scala
+++ b/src/main/scala/firrtl/transforms/OptimizationAnnotations.scala
@@ -8,11 +8,33 @@ import firrtl.passes.PassException
 /** Indicate that DCE should not be run */
 case object NoDCEAnnotation extends NoTargetAnnotation
 
+/** Lets an annotation mark its ReferenceTarget members as DontTouch
+  *
+  * This permits a transform to run and remove its associated annotations,
+  * thus making their ReferenceTargets new candidates for optimization. This
+  * removes the need for the pass writer to reason about pre-existing
+  * DontTouchAnnotations that may touch the same node.
+  */
+trait HasDontTouches { self: Annotation =>
+  def dontTouches: Iterable[ReferenceTarget]
+}
+
+/**
+  * A globalized form of HasDontTouches which applies to all ReferenceTargets
+  * provided with the annotation
+  */
+trait DontTouchAllTargets extends HasDontTouches { self: Annotation =>
+  def dontTouches: Iterable[ReferenceTarget] = getTargets.collect {
+    case rT: ReferenceTarget => rT
+  }
+}
+
 /** A component that should be preserved
   *
   * DCE treats the component as a top-level sink of the circuit
   */
-case class DontTouchAnnotation(target: ReferenceTarget) extends SingleTargetAnnotation[ReferenceTarget] {
+case class DontTouchAnnotation(target: ReferenceTarget)
+    extends SingleTargetAnnotation[ReferenceTarget] with DontTouchAllTargets {
   def targets = Seq(target)
   def duplicate(n: ReferenceTarget) = this.copy(n)
 }
diff --git a/src/test/scala/firrtlTests/ConstantPropagationTests.scala b/src/test/scala/firrtlTests/ConstantPropagationTests.scala
index 18579ca8..189809a6 100644
--- a/src/test/scala/firrtlTests/ConstantPropagationTests.scala
+++ b/src/test/scala/firrtlTests/ConstantPropagationTests.scala
@@ -792,7 +792,20 @@ class ConstantPropagationIntegrationSpec extends LowTransformSpec {
     execute(input, check, Seq(dontTouch("Top.z")))
   }
 
-  "ConstProp" should "NOT optimize across dontTouch on registers" in {
+  it should "NOT optimize across nodes marked dontTouch by other annotations" in {
+      val input =
+        """circuit Top :
+          |  module Top :
+          |    input x : UInt<1>
+          |    output y : UInt<1>
+          |    node z = x
+          |    y <= z""".stripMargin
+      val check = input
+      val dontTouchRT = annotations.ModuleTarget("Top", "Top").ref("z")
+    execute(input, check, Seq(AnnotationWithDontTouches(dontTouchRT)))
+  }
+
+  it should "NOT optimize across dontTouch on registers" in {
       val input =
         """circuit Top :
           |  module Top :
diff --git a/src/test/scala/firrtlTests/DCETests.scala b/src/test/scala/firrtlTests/DCETests.scala
index a9dbdda2..bfd47042 100644
--- a/src/test/scala/firrtlTests/DCETests.scala
+++ b/src/test/scala/firrtlTests/DCETests.scala
@@ -11,6 +11,13 @@ import firrtl.passes.memlib.SimpleTransform
 import java.io.File
 import java.nio.file.Paths
 
+case class AnnotationWithDontTouches(target: ReferenceTarget)
+    extends SingleTargetAnnotation[ReferenceTarget] with HasDontTouches {
+  def targets = Seq(target)
+  def duplicate(n: ReferenceTarget) = this.copy(n)
+  def dontTouches: Seq[ReferenceTarget] = targets
+}
+
 class DCETests extends FirrtlFlatSpec {
   // Not using executeTest because it is for positive testing, we need to check that stuff got
   // deleted
@@ -63,6 +70,24 @@ class DCETests extends FirrtlFlatSpec {
         |    z <= x""".stripMargin
     exec(input, check, Seq(dontTouch("Top.a")))
   }
+  "Unread wire marked dont touch by another annotation" should "NOT be deleted" in {
+    val input =
+      """circuit Top :
+        |  module Top :
+        |    input x : UInt<1>
+        |    output z : UInt<1>
+        |    wire a : UInt<1>
+        |    z <= x
+        |    a <= x""".stripMargin
+    val check =
+      """circuit Top :
+        |  module Top :
+        |    input x : UInt<1>
+        |    output z : UInt<1>
+        |    node a = x
+        |    z <= x""".stripMargin
+    exec(input, check, Seq(AnnotationWithDontTouches(ModuleTarget("Top", "Top").ref("a"))))
+  }
   "Unread register" should "be deleted" in {
     val input =
       """circuit Top :