5 files changed, 50 insertions, 16 deletions

diff --git a/build.sbt b/build.sbt
index 9fba6a6b..882ac0fd 100644
--- a/build.sbt
+++ b/build.sbt
@@ -89,7 +89,7 @@ lazy val testAssemblySettings = Seq(
 
 lazy val antlrSettings = Seq(
   Antlr4 / antlr4GenVisitor := true,
-  Antlr4 / antlr4GenListener := false,
+  Antlr4 / antlr4GenListener := true,
   Antlr4 / antlr4PackageName := Option("firrtl.antlr"),
   Antlr4 / antlr4Version := "4.9.3",
   Antlr4 / javaSource := (Compile / sourceManaged).value
diff --git a/build.sc b/build.sc
index 59e6d4c0..de5922d8 100644
--- a/build.sc
+++ b/build.sc
@@ -139,7 +139,7 @@ class firrtlCrossModule(val crossScalaVersion: String)
           antlrSource().path.toString,
           "-package",
           "firrtl.antlr",
-          "-no-listener",
+          "-listener",
           "-visitor",
           antlrSource().path.toString
         ).call()
@@ -152,7 +152,7 @@ class firrtlCrossModule(val crossScalaVersion: String)
           antlrSource().path.toString,
           "-package",
           "firrtl.antlr",
-          "-no-listener",
+          "-listener",
           "-visitor",
           antlrSource().path.toString
         ).call()
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)
+  }
+}