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>

3 files changed, 32 insertions, 5 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)
 }