Dedup attribute annos (#2297)

* Add new util "groupByIntoSeq"
* Restore annotation order when dedupping annotations
* Attribute annotations now deduplicate
* Implement doc string anno dedup

Co-authored-by: Jack Koenig <koenig@sifive.com>

5 files changed, 344 insertions, 8 deletions

diff --git a/src/main/scala/firrtl/AddDescriptionNodes.scala b/src/main/scala/firrtl/AddDescriptionNodes.scala
index 123ae6e3..90e38beb 100644
--- a/src/main/scala/firrtl/AddDescriptionNodes.scala
+++ b/src/main/scala/firrtl/AddDescriptionNodes.scala
@@ -28,6 +28,13 @@ case class DocStringAnnotation(target: Target, description: String) extends Desc
       case Some(seq) => seq.map(n => this.copy(target = n))
     }
   }
+  override private[firrtl] def dedup: Option[(Any, Annotation, ReferenceTarget)] = this match {
+    case a @ DocStringAnnotation(refTarget: ReferenceTarget, _) =>
+      Some(((refTarget.pathlessTarget, description), copy(target = refTarget.pathlessTarget), refTarget))
+    case a @ DocStringAnnotation(pathTarget: InstanceTarget, _) =>
+      Some(((pathTarget.pathlessTarget, description), copy(target = pathTarget.pathlessTarget), pathTarget.asReference))
+    case _ => None
+  }
 }
 
 /**
@@ -42,6 +49,13 @@ case class AttributeAnnotation(target: Target, description: String) extends Desc
       case Some(seq) => seq.map(n => this.copy(target = n))
     }
   }
+  override private[firrtl] def dedup: Option[(Any, Annotation, ReferenceTarget)] = this match {
+    case a @ AttributeAnnotation(refTarget: ReferenceTarget, _) =>
+      Some(((refTarget.pathlessTarget, description), copy(target = refTarget.pathlessTarget), refTarget))
+    case a @ AttributeAnnotation(pathTarget: InstanceTarget, _) =>
+      Some(((pathTarget.pathlessTarget, description), copy(target = pathTarget.pathlessTarget), pathTarget.asReference))
+    case _ => None
+  }
 }
 
 /**
diff --git a/src/main/scala/firrtl/Utils.scala b/src/main/scala/firrtl/Utils.scala
index e29c1a3b..e2bb06ff 100644
--- a/src/main/scala/firrtl/Utils.scala
+++ b/src/main/scala/firrtl/Utils.scala
@@ -967,6 +967,17 @@ object Utils extends LazyLogging {
     Mux(cond, tval, fval, tval.tpe)
   }
 
+  /** Similar to Seq.groupBy except that it preserves ordering of elements within each group */
+  def groupByIntoSeq[A, K](xs: Iterable[A])(f: A => K): Seq[(K, Seq[A])] = {
+    val map = mutable.LinkedHashMap.empty[K, mutable.ListBuffer[A]]
+    for (x <- xs) {
+      val key = f(x)
+      val l = map.getOrElseUpdate(key, mutable.ListBuffer.empty[A])
+      l += x
+    }
+    map.view.map({ case (k, vs) => k -> vs.toList }).toList
+  }
+
   object True {
     private val _True = UIntLiteral(1, IntWidth(1))
 
diff --git a/src/main/scala/firrtl/backends/verilog/VerilogEmitter.scala b/src/main/scala/firrtl/backends/verilog/VerilogEmitter.scala
index 8b20d365..4f62f27e 100644
--- a/src/main/scala/firrtl/backends/verilog/VerilogEmitter.scala
+++ b/src/main/scala/firrtl/backends/verilog/VerilogEmitter.scala
@@ -8,6 +8,7 @@ import firrtl.Utils._
 import firrtl.WrappedExpression._
 import firrtl.traversals.Foreachers._
 import firrtl.annotations.{
+  Annotation,
   CircuitTarget,
   MemoryInitAnnotation,
   MemoryLoadFileType,
@@ -509,14 +510,14 @@ class VerilogEmitter extends SeqTransform with Emitter {
       case m: SingleTargetAnnotation[ReferenceTarget] @unchecked with EmissionOption => m
     }
 
-    // Check for non-local memory annotations (error if found)
-    emissionAnnos.foreach {
-      case a: MemoryInitAnnotation => {
-        if (!a.target.isLocal)
+    annotations.foreach {
+      case a: Annotation if a.dedup.nonEmpty =>
+        val (_, _, target) = a.dedup.get
+        if (!target.isLocal) {
           throw new FirrtlUserException(
-            "At least one memory annotation did not deduplicate: got non-local annotation $a from [[DedupAnnotationsTransform]]"
+            "At least one dedupable annotation did not deduplicate: got non-local annotation $a from [[DedupAnnotationsTransform]]"
           )
-      }
+        }
       case _ =>
     }
 
diff --git a/src/main/scala/firrtl/transforms/DedupAnnotations.scala b/src/main/scala/firrtl/transforms/DedupAnnotations.scala
index 9355b5c3..cad4d2be 100644
--- a/src/main/scala/firrtl/transforms/DedupAnnotations.scala
+++ b/src/main/scala/firrtl/transforms/DedupAnnotations.scala
@@ -6,7 +6,7 @@ package transforms
 import firrtl.ir._
 import firrtl.Mappers._
 import firrtl.options.Dependency
-import firrtl.Utils.BoolType
+import firrtl.Utils.{groupByIntoSeq, BoolType}
 import firrtl.annotations.Annotation
 import scala.collection.mutable.Buffer
 import firrtl.annotations.MemoryFileInlineAnnotation
@@ -61,7 +61,7 @@ object DedupAnnotationsTransform {
     }
 
     // Partition the dedupable annotations into groups that *should* deduplicate into the same annotation
-    val shouldDedup: Map[Any, ArrayBuffer[DedupableRepr]] = canDedup.groupBy(_.dedupKey)
+    val shouldDedup: Seq[(Any, Seq[DedupableRepr])] = groupByIntoSeq(canDedup)(_.dedupKey)
     shouldDedup.foreach {
       case ((target: ReferenceTarget, _), dedupableAnnos) =>
         val originalAnnos = dedupableAnnos.map(_.original)
diff --git a/src/test/scala/firrtlTests/VerilogEmitterTests.scala b/src/test/scala/firrtlTests/VerilogEmitterTests.scala
index ec30e55c..f4d13ac6 100644
--- a/src/test/scala/firrtlTests/VerilogEmitterTests.scala
+++ b/src/test/scala/firrtlTests/VerilogEmitterTests.scala
@@ -1004,6 +1004,316 @@ class VerilogDescriptionEmitterSpec extends FirrtlFlatSpec {
       assert(output.contains(c))
     }
   }
+
+  "Attribute annotations" should "be deduplicated" in {
+    val compiler = new VerilogCompiler
+    val input =
+      """
+        |circuit Top:
+        |  module Child:
+        |    input clock : Clock
+        |    input rAddr : UInt<5>
+        |    input rEnable : UInt<1>
+        |    input wAddr : UInt<5>
+        |    input wData : UInt<8>
+        |    input wEnable : UInt<1>
+        |    output rData : UInt<8>
+        |
+        |    mem m:
+        |      data-type => UInt<8>
+        |      depth => 32
+        |      reader => r
+        |      writer => w
+        |      read-latency => 1
+        |      write-latency => 1
+        |      read-under-write => new
+        |
+        |    m.r.clk <= clock
+        |    m.r.addr <= rAddr
+        |    m.r.en <= rEnable
+        |    rData <= m.r.data
+        |
+        |    m.w.clk <= clock
+        |    m.w.addr <= wAddr
+        |    m.w.en <= wEnable
+        |    m.w.data <= wData
+        |    m.w.mask is invalid
+        |
+        |  module Top:
+        |    input clock : Clock
+        |    input rAddr : UInt<5>
+        |    input rEnable : UInt<1>
+        |    input wAddr : UInt<5>
+        |    input wData : UInt<8>
+        |    input wEnable : UInt<1>
+        |    output rData : UInt<8>[2]
+        |
+        |    inst c1 of Child
+        |    c1.clock <= clock
+        |    c1.rAddr <= rAddr
+        |    c1.rEnable <= rEnable
+        |    c1.wAddr <= wAddr
+        |    c1.wData <= wData
+        |    c1.wEnable <= wEnable
+        |
+        |    inst c2 of Child
+        |    c2.clock <= clock
+        |    c2.rAddr <= rAddr
+        |    c2.rEnable <= rEnable
+        |    c2.wAddr <= wAddr
+        |    c2.wData <= wData
+        |    c2.wEnable <= wEnable
+        |
+        |    rData[0] <= c1.rData
+        |    rData[1] <= c2.rData
+        |""".stripMargin
+    val desc = "child module"
+    val child1MRef = CircuitTarget("Top").module("Top").instOf("c1", "Child").ref("m")
+    val child2MRef = CircuitTarget("Top").module("Top").instOf("c2", "Child").ref("m")
+    val annos = Seq(
+      AttributeAnnotation(child1MRef, desc),
+      AttributeAnnotation(child2MRef, desc)
+    )
+
+    val finalState = compiler.compileAndEmit(CircuitState(parse(input), ChirrtlForm, annos), Seq.empty)
+    val output = finalState.getEmittedCircuit.value
+    assert(output.contains(s"  (* $desc *)"))
+  }
+
+  "Doc string annotations" should "be deduplicated" in {
+    val compiler = new VerilogCompiler
+    val input =
+      """
+        |circuit Top:
+        |  module Child:
+        |    input clock : Clock
+        |    input rAddr : UInt<5>
+        |    input rEnable : UInt<1>
+        |    input wAddr : UInt<5>
+        |    input wData : UInt<8>
+        |    input wEnable : UInt<1>
+        |    output rData : UInt<8>
+        |
+        |    mem m:
+        |      data-type => UInt<8>
+        |      depth => 32
+        |      reader => r
+        |      writer => w
+        |      read-latency => 1
+        |      write-latency => 1
+        |      read-under-write => new
+        |
+        |    m.r.clk <= clock
+        |    m.r.addr <= rAddr
+        |    m.r.en <= rEnable
+        |    rData <= m.r.data
+        |
+        |    m.w.clk <= clock
+        |    m.w.addr <= wAddr
+        |    m.w.en <= wEnable
+        |    m.w.data <= wData
+        |    m.w.mask is invalid
+        |
+        |  module Top:
+        |    input clock : Clock
+        |    input rAddr : UInt<5>
+        |    input rEnable : UInt<1>
+        |    input wAddr : UInt<5>
+        |    input wData : UInt<8>
+        |    input wEnable : UInt<1>
+        |    output rData : UInt<8>[2]
+        |
+        |    inst c1 of Child
+        |    c1.clock <= clock
+        |    c1.rAddr <= rAddr
+        |    c1.rEnable <= rEnable
+        |    c1.wAddr <= wAddr
+        |    c1.wData <= wData
+        |    c1.wEnable <= wEnable
+        |
+        |    inst c2 of Child
+        |    c2.clock <= clock
+        |    c2.rAddr <= rAddr
+        |    c2.rEnable <= rEnable
+        |    c2.wAddr <= wAddr
+        |    c2.wData <= wData
+        |    c2.wEnable <= wEnable
+        |
+        |    rData[0] <= c1.rData
+        |    rData[1] <= c2.rData
+        |""".stripMargin
+    val check =
+      """  /* line1
+        |   *
+        |   * line2
+        |   */
+        |""".stripMargin
+
+    val child1MRef = CircuitTarget("Top").module("Top").instOf("c1", "Child").ref("m")
+    val child2MRef = CircuitTarget("Top").module("Top").instOf("c2", "Child").ref("m")
+    val annos = Seq(
+      DocStringAnnotation(child1MRef, "line1"),
+      DocStringAnnotation(child1MRef, "line2"),
+      DocStringAnnotation(child2MRef, "line1"),
+      DocStringAnnotation(child2MRef, "line2")
+    )
+
+    val finalState = compiler.compileAndEmit(CircuitState(parse(input), ChirrtlForm, annos), Seq.empty)
+    val output = finalState.getEmittedCircuit.value
+    assert(output.contains(check))
+  }
+
+  "AttributeAnnotation" should "fail dedup if not all instances have the annotation" in {
+    val compiler = new VerilogCompiler
+    val input =
+      """
+        |circuit Top:
+        |  module Child:
+        |    input clock : Clock
+        |    input rAddr : UInt<5>
+        |    input rEnable : UInt<1>
+        |    input wAddr : UInt<5>
+        |    input wData : UInt<8>
+        |    input wEnable : UInt<1>
+        |    output rData : UInt<8>
+        |
+        |    mem m:
+        |      data-type => UInt<8>
+        |      depth => 32
+        |      reader => r
+        |      writer => w
+        |      read-latency => 1
+        |      write-latency => 1
+        |      read-under-write => new
+        |
+        |    m.r.clk <= clock
+        |    m.r.addr <= rAddr
+        |    m.r.en <= rEnable
+        |    rData <= m.r.data
+        |
+        |    m.w.clk <= clock
+        |    m.w.addr <= wAddr
+        |    m.w.en <= wEnable
+        |    m.w.data <= wData
+        |    m.w.mask is invalid
+        |
+        |  module Top:
+        |    input clock : Clock
+        |    input rAddr : UInt<5>
+        |    input rEnable : UInt<1>
+        |    input wAddr : UInt<5>
+        |    input wData : UInt<8>
+        |    input wEnable : UInt<1>
+        |    output rData : UInt<8>[2]
+        |
+        |    inst c1 of Child
+        |    c1.clock <= clock
+        |    c1.rAddr <= rAddr
+        |    c1.rEnable <= rEnable
+        |    c1.wAddr <= wAddr
+        |    c1.wData <= wData
+        |    c1.wEnable <= wEnable
+        |
+        |    inst c2 of Child
+        |    c2.clock <= clock
+        |    c2.rAddr <= rAddr
+        |    c2.rEnable <= rEnable
+        |    c2.wAddr <= wAddr
+        |    c2.wData <= wData
+        |    c2.wEnable <= wEnable
+        |
+        |    rData[0] <= c1.rData
+        |    rData[1] <= c2.rData
+        |""".stripMargin
+    val desc = "child module"
+    val child1MRef = CircuitTarget("Top").module("Top").instOf("c1", "Child").ref("m")
+    val child2MRef = CircuitTarget("Top").module("Top").instOf("c2", "Child").ref("m")
+    val annos = Seq(
+      AttributeAnnotation(child1MRef, desc)
+    )
+
+    assertThrows[FirrtlUserException] {
+      val finalState = compiler.compileAndEmit(CircuitState(parse(input), ChirrtlForm, annos), Seq.empty)
+      val output = finalState.getEmittedCircuit.value
+    }
+  }
+
+  "DocStringAnnotation" should "fail dedup if not all instances share the annotation" in {
+    val compiler = new VerilogCompiler
+    val input =
+      """
+        |circuit Top:
+        |  module Child:
+        |    input clock : Clock
+        |    input rAddr : UInt<5>
+        |    input rEnable : UInt<1>
+        |    input wAddr : UInt<5>
+        |    input wData : UInt<8>
+        |    input wEnable : UInt<1>
+        |    output rData : UInt<8>
+        |
+        |    mem m:
+        |      data-type => UInt<8>
+        |      depth => 32
+        |      reader => r
+        |      writer => w
+        |      read-latency => 1
+        |      write-latency => 1
+        |      read-under-write => new
+        |
+        |    m.r.clk <= clock
+        |    m.r.addr <= rAddr
+        |    m.r.en <= rEnable
+        |    rData <= m.r.data
+        |
+        |    m.w.clk <= clock
+        |    m.w.addr <= wAddr
+        |    m.w.en <= wEnable
+        |    m.w.data <= wData
+        |    m.w.mask is invalid
+        |
+        |  module Top:
+        |    input clock : Clock
+        |    input rAddr : UInt<5>
+        |    input rEnable : UInt<1>
+        |    input wAddr : UInt<5>
+        |    input wData : UInt<8>
+        |    input wEnable : UInt<1>
+        |    output rData : UInt<8>[2]
+        |
+        |    inst c1 of Child
+        |    c1.clock <= clock
+        |    c1.rAddr <= rAddr
+        |    c1.rEnable <= rEnable
+        |    c1.wAddr <= wAddr
+        |    c1.wData <= wData
+        |    c1.wEnable <= wEnable
+        |
+        |    inst c2 of Child
+        |    c2.clock <= clock
+        |    c2.rAddr <= rAddr
+        |    c2.rEnable <= rEnable
+        |    c2.wAddr <= wAddr
+        |    c2.wData <= wData
+        |    c2.wEnable <= wEnable
+        |
+        |    rData[0] <= c1.rData
+        |    rData[1] <= c2.rData
+        |""".stripMargin
+    val desc = "child module"
+    val child1MRef = CircuitTarget("Top").module("Top").instOf("c1", "Child").ref("m")
+    val child2MRef = CircuitTarget("Top").module("Top").instOf("c2", "Child").ref("m")
+    val annos = Seq(
+      DocStringAnnotation(child1MRef, "line1"),
+      DocStringAnnotation(child2MRef, "line2"),
+      DocStringAnnotation(child1MRef, "line1")
+    )
+
+    assertThrows[FirrtlUserException] {
+      val finalState = compiler.compileAndEmit(CircuitState(parse(input), ChirrtlForm, annos), Seq.empty)
+      val output = finalState.getEmittedCircuit.value
+    }
+  }
 }
 
 class EmittedMacroSpec extends FirrtlPropSpec {