Optimize Annotation.getTargets (#2244)

2 files changed, 47 insertions, 5 deletions

diff --git a/src/main/scala/firrtl/annotations/Annotation.scala b/src/main/scala/firrtl/annotations/Annotation.scala
index 8722c4fa..08555a84 100644
--- a/src/main/scala/firrtl/annotations/Annotation.scala
+++ b/src/main/scala/firrtl/annotations/Annotation.scala
@@ -5,6 +5,8 @@ package annotations
 
 import firrtl.options.StageUtils
 
+import scala.collection.Traversable
+
 case class AnnotationException(message: String) extends Exception(message)
 
 /** Base type of auxiliary information */
@@ -23,18 +25,19 @@ trait Annotation extends Product {
     * @param ls
     * @return
     */
-  private def extractComponents(ls: scala.collection.Traversable[_]): Seq[Target] = {
-    ls.collect {
+  private def extractComponents(ls: Traversable[_]): Traversable[Target] = {
+    ls.flatMap {
       case c: Target                          => Seq(c)
-      case o: Product                         => extractComponents(o.productIterator.toIterable)
       case x: scala.collection.Traversable[_] => extractComponents(x)
-    }.foldRight(Seq.empty[Target])((seq, c) => c ++ seq)
+      case o: Product                         => extractComponents(o.productIterator.toIterable)
+      case _ => Seq()
+    }
   }
 
   /** Returns all [[firrtl.annotations.Target Target]] members in this annotation
     * @return
     */
-  def getTargets: Seq[Target] = extractComponents(productIterator.toSeq)
+  def getTargets: Seq[Target] = extractComponents(productIterator.toIterable).toSeq
 }
 
 /** If an Annotation does not target any [[Named]] thing in the circuit, then all updates just
@@ -42,6 +45,8 @@ trait Annotation extends Product {
   */
 trait NoTargetAnnotation extends Annotation {
   def update(renames: RenameMap): Seq[NoTargetAnnotation] = Seq(this)
+
+  override def getTargets: Seq[Target] = Seq.empty
 }
 
 /** An Annotation that targets a single [[Named]] thing */
@@ -103,6 +108,8 @@ trait MultiTargetAnnotation extends Annotation {
     */
   def targets: Seq[Seq[Target]]
 
+  override def getTargets: Seq[Target] = targets.flatten
+
   /** Create another instance of this Annotation
     *
     * The inner Seqs correspond to the renames of the inner Seqs of targets
diff --git a/src/test/scala/firrtlTests/annotationTests/AnnotationSpec.scala b/src/test/scala/firrtlTests/annotationTests/AnnotationSpec.scala
new file mode 100644
index 00000000..2729b4dc
--- /dev/null
+++ b/src/test/scala/firrtlTests/annotationTests/AnnotationSpec.scala
@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package firrtlTests.annotationTests
+
+import firrtl.RenameMap
+import firrtl.annotations._
+import firrtl.testutils.FirrtlFlatSpec
+
+object AnnotationSpec {
+  case class TestAnno(pairs: List[(String, ReferenceTarget)]) extends Annotation {
+    def update(renames: RenameMap): Seq[Annotation] = {
+      val pairsx = pairs.flatMap {
+        case (n, t) =>
+          val ts = renames
+            .get(t)
+            .map(_.map(_.asInstanceOf[ReferenceTarget]))
+            .getOrElse(Seq(t))
+          ts.map(n -> _)
+      }
+      Seq(TestAnno(pairsx))
+    }
+  }
+}
+
+class AnnotationSpec extends FirrtlFlatSpec {
+  import AnnotationSpec._
+
+  behavior.of("Annotation.getTargets")
+
+  it should "not stack overflow" in {
+    val ref = CircuitTarget("Top").module("Foo").ref("vec")
+    val anno = TestAnno((0 until 10000).map(i => (i.toString, ref.index(i))).toList)
+    anno.getTargets should be(anno.pairs.map(_._2))
+  }
+}