From 17250fba841ae3129dc798c0bc48d10200be18ae Mon Sep 17 00:00:00 2001
From: Jack Koenig
Date: Tue, 31 Aug 2021 14:10:29 -0700
Subject: Use ANTLR Listener to save memory during parsing

The ANTLR-generated concrete syntax tree (CST) takes up much more memory
than the parsed .fir file. By using a Listener, we can construct the
FIRRTL AST live with CST construction and null out the CST as we consume
pieces of it. Not only does this improve performance, it drastically
reduces max memory use for the parser.
---
 src/main/scala/firrtl/Parser.scala          | 18 ++++++--------
 src/main/scala/firrtl/Visitor.scala         |  4 +--
 src/main/scala/firrtl/parser/Listener.scala | 38 +++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 13 deletions(-)
 create mode 100644 src/main/scala/firrtl/parser/Listener.scala

(limited to 'src')

diff --git a/src/main/scala/firrtl/Parser.scala b/src/main/scala/firrtl/Parser.scala
index 2d2bd350..bb93511c 100644
--- a/src/main/scala/firrtl/Parser.scala
+++ b/src/main/scala/firrtl/Parser.scala
@@ -6,6 +6,7 @@ import org.antlr.v4.runtime._
 import org.antlr.v4.runtime.atn._
 import logger.LazyLogging
 import firrtl.ir._
+import firrtl.parser.Listener
 import firrtl.Utils.time
 import firrtl.antlr.{FIRRTLParser, _}
 
@@ -29,29 +30,24 @@ object Parser extends LazyLogging {
 
   /** Parses a org.antlr.v4.runtime.CharStream and returns a parsed [[firrtl.ir.Circuit Circuit]] */
   def parseCharStream(charStream: CharStream, infoMode: InfoMode): Circuit = {
-    val (parseTimeMillis, cst) = time {
+    val (parseTimeMillis, ast) = time {
       val parser = {
         val lexer = new FIRRTLLexer(charStream)
         new FIRRTLParser(new CommonTokenStream(lexer))
       }
 
+      val listener = new Listener(infoMode)
+
       parser.getInterpreter.setPredictionMode(PredictionMode.SLL)
+      parser.addParseListener(listener)
 
       // Concrete Syntax Tree
-      val cst = parser.circuit
+      parser.circuit
 
       val numSyntaxErrors = parser.getNumberOfSyntaxErrors
       if (numSyntaxErrors > 0) throw new SyntaxErrorsException(s"$numSyntaxErrors syntax error(s) detected")
-      cst
-    }
 
-    val visitor = new Visitor(infoMode)
-    val (visitTimeMillis, visit) = time {
-      visitor.visit(cst)
-    }
-    val ast = visit match {
-      case c: Circuit => c
-      case x => throw new ClassCastException("Error! AST not rooted with Circuit node!")
+      listener.getCircuit
     }
 
     ast
diff --git a/src/main/scala/firrtl/Visitor.scala b/src/main/scala/firrtl/Visitor.scala
index b2c4896f..f3b6837e 100644
--- a/src/main/scala/firrtl/Visitor.scala
+++ b/src/main/scala/firrtl/Visitor.scala
@@ -58,7 +58,7 @@ class Visitor(infoMode: InfoMode) extends AbstractParseTreeVisitor[FirrtlNode] w
 
   private def string2Int(s: String): Int = string2BigInt(s).toInt
 
-  private def visitInfo(ctx: Option[InfoContext], parentCtx: ParserRuleContext): Info = {
+  private[firrtl] def visitInfo(ctx: Option[InfoContext], parentCtx: ParserRuleContext): Info = {
     // Convert a compressed FileInfo string into either into a singular FileInfo or a MultiInfo
     // consisting of several FileInfos
     def parseCompressedInfo(escaped: String): Info = {
@@ -130,7 +130,7 @@ class Visitor(infoMode: InfoMode) extends AbstractParseTreeVisitor[FirrtlNode] w
   private def visitCircuit(ctx: CircuitContext): Circuit =
     Circuit(visitInfo(Option(ctx.info), ctx), ctx.module.asScala.map(visitModule).toSeq, ctx.id.getText)
 
-  private def visitModule(ctx: ModuleContext): DefModule = {
+  private[firrtl] def visitModule(ctx: ModuleContext): DefModule = {
     val info = visitInfo(Option(ctx.info), ctx)
     ctx.getChild(0).getText match {
       case "module" =>
diff --git a/src/main/scala/firrtl/parser/Listener.scala b/src/main/scala/firrtl/parser/Listener.scala
new file mode 100644
index 00000000..ffaa22b2
--- /dev/null
+++ b/src/main/scala/firrtl/parser/Listener.scala
@@ -0,0 +1,38 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package firrtl.parser
+
+import firrtl.antlr.{FIRRTLParser, _}
+import firrtl.Visitor
+import firrtl.Parser.InfoMode
+import firrtl.ir._
+
+import scala.collection.mutable
+import scala.concurrent.{Await, Future}
+import scala.concurrent.duration.Duration
+
+private[firrtl] class Listener(infoMode: InfoMode) extends FIRRTLBaseListener {
+  private var main: Option[String] = None
+  private var info: Option[Info] = None
+  private val modules = mutable.ArrayBuffer.empty[DefModule]
+
+  private val visitor = new Visitor(infoMode)
+
+  override def exitModule(ctx: FIRRTLParser.ModuleContext): Unit = {
+    val m = visitor.visitModule(ctx)
+    ctx.children = null // Null out to save memory
+    modules += m
+  }
+
+  override def exitCircuit(ctx: FIRRTLParser.CircuitContext): Unit = {
+    info = Some(visitor.visitInfo(Option(ctx.info), ctx))
+    main = Some(ctx.id.getText)
+    ctx.children = null // Null out to save memory
+  }
+
+  def getCircuit: Circuit = {
+    require(main.nonEmpty)
+    val mods = modules.toSeq
+    Circuit(info.get, mods, main.get)
+  }
+}
-- 
cgit v1.2.3