3 files changed, 47 insertions, 13 deletions
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)
+  }
+}