Improve Parser and Visitor (#819)

* Update Parser to use ANTLR CharStreams

This removes some unnecessary object creation in String reading and
manipulation

* Remove two unnecessary traversals from Block construction in Visitor

3 files changed, 20 insertions, 16 deletions

diff --git a/src/main/scala/firrtl/Driver.scala b/src/main/scala/firrtl/Driver.scala
index 72128274..ddc06d2a 100644
--- a/src/main/scala/firrtl/Driver.scala
+++ b/src/main/scala/firrtl/Driver.scala
@@ -169,7 +169,7 @@ object Driver {
         throw new OptionsException(msg)
       }
       firrtlConfig.firrtlCircuit.getOrElse {
-        val source = firrtlConfig.firrtlSource.map(_.split("\n").toIterator).getOrElse {
+        firrtlConfig.firrtlSource.map(x => Parser.parseString(x, firrtlConfig.infoMode)).getOrElse {
           if (optionsManager.topName.isEmpty && firrtlConfig.inputFileNameOverride.isEmpty) {
             val message = "either top-name or input-file-override must be set"
             throw new OptionsException(message)
@@ -183,7 +183,7 @@ object Driver {
           }
           val inputFileName = firrtlConfig.getInputFileName(optionsManager)
           try {
-            io.Source.fromFile(inputFileName).getLines()
+            Parser.parseFile(inputFileName, firrtlConfig.infoMode)
           }
           catch {
             case _: FileNotFoundException =>
@@ -191,7 +191,6 @@ object Driver {
               throw new OptionsException(message)
           }
         }
-        Parser.parse(source, firrtlConfig.infoMode)
       }
     }
   }
diff --git a/src/main/scala/firrtl/Parser.scala b/src/main/scala/firrtl/Parser.scala
index 70ae7739..8b0f5b6b 100644
--- a/src/main/scala/firrtl/Parser.scala
+++ b/src/main/scala/firrtl/Parser.scala
@@ -2,8 +2,6 @@
 
 package firrtl
 
-import java.io.{ByteArrayInputStream, SequenceInputStream}
-
 import org.antlr.v4.runtime._
 import org.antlr.v4.runtime.atn._
 import com.typesafe.scalalogging.LazyLogging
@@ -21,16 +19,20 @@ case class SyntaxErrorsException(message: String) extends ParserException(messag
 
 
 object Parser extends LazyLogging {
-  /** Takes Iterator over lines of FIRRTL, returns FirrtlNode (root node is Circuit) */
-  def parse(lines: Iterator[String], infoMode: InfoMode = UseInfo): Circuit = {
 
+  /** Parses a file in a given filename and returns a parsed [[Circuit]] */
+  def parseFile(filename: String, infoMode: InfoMode): Circuit =
+    parseCharStream(CharStreams.fromFileName(filename), infoMode)
+
+  /** Parses a String and returns a parsed [[Circuit]] */
+  def parseString(text: String, infoMode: InfoMode): Circuit =
+    parseCharStream(CharStreams.fromString(text), infoMode)
+
+  /** Parses a org.antlr.v4.runtime.CharStream and returns a parsed [[Circuit]] */
+  def parseCharStream(charStream: CharStream, infoMode: InfoMode): Circuit = {
     val (parseTimeMillis, cst) = time {
       val parser = {
-        import scala.collection.JavaConverters._
-        val inStream = new SequenceInputStream(
-          lines.map{s => new ByteArrayInputStream((s + "\n").getBytes("UTF-8")) }.asJavaEnumeration
-        )
-        val lexer = new FIRRTLLexer(new ANTLRInputStream(inStream))
+        val lexer = new FIRRTLLexer(charStream)
         new FIRRTLParser(new CommonTokenStream(lexer))
       }
 
@@ -55,10 +57,13 @@ object Parser extends LazyLogging {
 
     ast
   }
+  /** Takes Iterator over lines of FIRRTL, returns FirrtlNode (root node is Circuit) */
+  def parse(lines: Iterator[String], infoMode: InfoMode = UseInfo): Circuit =
+    parseString(lines.mkString("\n"), infoMode)
 
-  def parse(lines: Seq[String]): Circuit = parse(lines.iterator)
+  def parse(lines: Seq[String]): Circuit = parseString(lines.mkString("\n"), UseInfo)
 
-  def parse(text: String): Circuit = parse(text split "\n")
+  def parse(text: String): Circuit = parseString(text, UseInfo)
 
   sealed abstract class InfoMode
 
diff --git a/src/main/scala/firrtl/Visitor.scala b/src/main/scala/firrtl/Visitor.scala
index f47d7ee3..7b965bff 100644
--- a/src/main/scala/firrtl/Visitor.scala
+++ b/src/main/scala/firrtl/Visitor.scala
@@ -144,10 +144,10 @@ class Visitor(infoMode: InfoMode) extends FIRRTLBaseVisitor[FirrtlNode] {
   }
 
   private def visitBlock[FirrtlNode](ctx: FIRRTLParser.ModuleBlockContext): Statement =
-    Block(ctx.simple_stmt().asScala.map(_.stmt).filter(_ != null).map(visitStmt))
+    Block(ctx.simple_stmt().asScala.flatMap(x => Option(x.stmt).map(visitStmt)))
 
   private def visitSuite[FirrtlNode](ctx: FIRRTLParser.SuiteContext): Statement =
-    Block(ctx.simple_stmt().asScala.map(_.stmt).filter(_ != null).map(visitStmt))
+    Block(ctx.simple_stmt().asScala.flatMap(x => Option(x.stmt).map(visitStmt)))
 
 
   // Memories are fairly complicated to translate thus have a dedicated method