Performance fix of Uniquify for deep bundles (#980)

3 files changed, 43 insertions, 15 deletions

diff --git a/src/main/scala/firrtl/annotations/AnnotationUtils.scala b/src/main/scala/firrtl/annotations/AnnotationUtils.scala
index ba9220f7..72765ab7 100644
--- a/src/main/scala/firrtl/annotations/AnnotationUtils.scala
+++ b/src/main/scala/firrtl/annotations/AnnotationUtils.scala
@@ -51,8 +51,8 @@ object AnnotationUtils {
     case Some(_) =>
       val i = s.indexWhere(c => "[].".contains(c))
       s.slice(0, i) match {
-        case "" => Seq(s(i).toString) ++ tokenize(s.drop(i + 1))
-        case x => Seq(x, s(i).toString) ++ tokenize(s.drop(i + 1))
+        case "" => s(i).toString +: tokenize(s.drop(i + 1))
+        case x => x +: s(i).toString +: tokenize(s.drop(i + 1))
       }
     case None if s == "" => Nil
     case None => Seq(s)
diff --git a/src/main/scala/firrtl/passes/Uniquify.scala b/src/main/scala/firrtl/passes/Uniquify.scala
index 2b6fa55d..c13fa261 100644
--- a/src/main/scala/firrtl/passes/Uniquify.scala
+++ b/src/main/scala/firrtl/passes/Uniquify.scala
@@ -3,14 +3,16 @@
 package firrtl.passes
 
 import com.typesafe.scalalogging.LazyLogging
-import scala.annotation.tailrec
 
+import scala.annotation.tailrec
 import firrtl._
 import firrtl.ir._
 import firrtl.Utils._
 import firrtl.Mappers._
 import MemPortUtils.memType
 
+import scala.collection.mutable
+
 /** Resolve name collisions that would occur in [[LowerTypes]]
   *
   *  @note Must be run after [[InferTypes]] because [[ir.DefNode]]s need type
@@ -78,36 +80,43 @@ object Uniquify extends Transform {
       t: BundleType,
       namespace: collection.mutable.HashSet[String])
       (implicit sinfo: Info, mname: String): BundleType = {
-    def recUniquifyNames(t: Type, namespace: collection.mutable.HashSet[String]): Type = t match {
+    def recUniquifyNames(t: Type, namespace: collection.mutable.HashSet[String]): (Type, Seq[String]) = t match {
       case tx: BundleType =>
         // First add everything
-        val newFields = tx.fields map { f =>
+        val newFieldsAndElts = tx.fields map { f =>
           val newName = findValidPrefix(f.name, Seq(""), namespace)
           namespace += newName
           Field(newName, f.flip, f.tpe)
         } map { f => f.tpe match {
-          case _: GroundType => f
+          case _: GroundType => (f, Seq[String](f.name))
           case _ =>
-            val tpe = recUniquifyNames(f.tpe, collection.mutable.HashSet())
-            val elts = enumerateNames(tpe)
+            val (tpe, eltsx) = recUniquifyNames(f.tpe, collection.mutable.HashSet())
             // Need leading _ for findValidPrefix, it doesn't add _ for checks
-            val eltsNames = elts map (e => "_" + LowerTypes.loweredName(e))
+            val eltsNames: Seq[String] = eltsx map (e => "_" + e)
             val prefix = findValidPrefix(f.name, eltsNames, namespace)
             // We added f.name in previous map, delete if we change it
             if (prefix != f.name) {
               namespace -= f.name
               namespace += prefix
             }
-            namespace ++= (elts map (e => LowerTypes.loweredName(prefix +: e)))
-            Field(prefix, f.flip, tpe)
+            val newElts: Seq[String] = eltsx map (e => LowerTypes.loweredName(prefix +: Seq(e)))
+            namespace ++= newElts
+            (Field(prefix, f.flip, tpe), prefix +: newElts)
           }
         }
-        BundleType(newFields)
+        val (newFields, elts) = newFieldsAndElts.unzip
+        (BundleType(newFields), elts.flatten)
       case tx: VectorType =>
-        VectorType(recUniquifyNames(tx.tpe, namespace), tx.size)
-      case tx => tx
+        val (tpe, elts) = recUniquifyNames(tx.tpe, namespace)
+        val newElts = ((0 until tx.size) map (i => i.toString)) ++
+          ((0 until tx.size) flatMap { i =>
+            elts map (e => LowerTypes.loweredName(Seq(i.toString, e)))
+          })
+        (VectorType(tpe, tx.size), newElts)
+      case tx => (tx, Nil)
     }
-    recUniquifyNames(t, namespace) match {
+    val (tpe, _) = recUniquifyNames(t, namespace)
+    tpe match {
       case tx: BundleType => tx
       case tx => throwInternalError(s"uniquifyNames: shouldn't be here - $tx")
     }
diff --git a/src/test/scala/firrtlTests/UniquifySpec.scala b/src/test/scala/firrtlTests/UniquifySpec.scala
index 43d1e733..bf0586f3 100644
--- a/src/test/scala/firrtlTests/UniquifySpec.scala
+++ b/src/test/scala/firrtlTests/UniquifySpec.scala
@@ -283,4 +283,23 @@ class UniquifySpec extends FirrtlFlatSpec {
 
     executeTest(input, expected)
   }
+
+  it should "quickly rename deep bundles" in {
+    def mkType(i: Int): String = {
+      if(i == 0) "UInt<8>" else s"{x: ${mkType(i - 1)}}"
+    }
+
+    val depth = 500
+
+    val input =
+      s"""circuit Test:
+         |  module Test :
+         |    input in: ${mkType(depth)}
+         |    output out: ${mkType(depth)}
+         |    out <= in
+         |""".stripMargin
+
+    val (ms, _) = Utils.time(compileToVerilog(input))
+    (ms < 8000) shouldBe true
+  }
 }