From 57b8a395ee8d5fdabb2deed3db7d0c644f0a7eed Mon Sep 17 00:00:00 2001
From: Jack Koenig <koenig@sifive.com>
Date: Mon, 20 Mar 2023 17:28:59 -0700
Subject: [PATCH] Make return value of Serializer.lazily lazy (#2627)

Directly subclassing Iterable is lazy-ish, but if you call any operation
on the resulting value (eg. map or ++) it will evaluate the Iterable and
return a List.
---
 src/main/scala/firrtl/ir/Serializer.scala     |  2 +-
 .../scala/firrtlTests/SerializerSpec.scala    | 44 +++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/src/main/scala/firrtl/ir/Serializer.scala b/src/main/scala/firrtl/ir/Serializer.scala
index 3ef9a96bc1..cf919b374b 100644
--- a/src/main/scala/firrtl/ir/Serializer.scala
+++ b/src/main/scala/firrtl/ir/Serializer.scala
@@ -69,7 +69,7 @@ object Serializer {
       case n: Circuit   => sIt(n)(indent)
       case other => Iterator(serialize(other, indent))
     }
-  }
+  }.view // TODO replace .view with constructing a view directly above, but must drop 2.12 first.
 
   /** Converts a `Constraint` into its string representation. */
   def serialize(con: Constraint): String = {
diff --git a/src/test/scala/firrtlTests/SerializerSpec.scala b/src/test/scala/firrtlTests/SerializerSpec.scala
index 0e72b97fdc..0af05e6045 100644
--- a/src/test/scala/firrtlTests/SerializerSpec.scala
+++ b/src/test/scala/firrtlTests/SerializerSpec.scala
@@ -136,4 +136,48 @@ class SerializerSpec extends AnyFlatSpec with Matchers {
     val readNew = parseSerializeParse(SMemTestCircuit.src(" new"))
     assert(SMemTestCircuit.findRuw(readNew) == ReadUnderWrite.New)
   }
+
+  it should "support lazy serialization" in {
+    var stmtSerialized = false
+    case class HackStmt(stmt: Statement) extends Statement {
+      def serialize: String = {
+        stmtSerialized = true
+        stmt.serialize
+      }
+      def foreachExpr(f:   Expression => Unit):       Unit = stmt.foreachExpr(f)
+      def foreachInfo(f:   Info => Unit):             Unit = stmt.foreachInfo(f)
+      def foreachStmt(f:   Statement => Unit):        Unit = stmt.foreachStmt(f)
+      def foreachString(f: String => Unit):           Unit = stmt.foreachString(f)
+      def foreachType(f:   Type => Unit):             Unit = stmt.foreachType(f)
+      def mapExpr(f:       Expression => Expression): Statement = this.copy(stmt.mapExpr(f))
+      def mapInfo(f:       Info => Info):             Statement = this.copy(stmt.mapInfo(f))
+      def mapStmt(f:       Statement => Statement):   Statement = this.copy(stmt.mapStmt(f))
+      def mapString(f:     String => String):         Statement = this.copy(stmt.mapString(f))
+      def mapType(f:       Type => Type):             Statement = this.copy(stmt.mapType(f))
+    }
+
+    val stmt = HackStmt(DefNode(NoInfo, "foo", Reference("bar")))
+    val it: Iterable[String] = Serializer.lazily(stmt)
+    assert(!stmtSerialized, "We should be able to construct the serializer lazily")
+
+    var mapExecuted = false
+    val it2: Iterable[String] = it.map { x =>
+      mapExecuted = true
+      x + ","
+    }
+    assert(!stmtSerialized && !mapExecuted, "We should be able to map the serializer lazily")
+
+    var appendExecuted = false
+    val it3: Iterable[String] = it2 ++ Seq("hi").view.map { x =>
+      appendExecuted = true
+      x
+    }
+    assert(!stmtSerialized && !mapExecuted && !appendExecuted, "We should be able to append to the serializer lazily")
+
+    val result = it3.mkString
+    assert(
+      stmtSerialized && mapExecuted && appendExecuted,
+      "Once we traverse the serializer, everything should execute"
+    )
+  }
 }