Handle references better in ANTLR Parser

Tweak the grammar to handle references without left-recursion. Also
split references and subreferences out from the regular expression rule
to make their parsing more efficient.

2 files changed, 56 insertions, 34 deletions

diff --git a/src/main/antlr4/FIRRTL.g4 b/src/main/antlr4/FIRRTL.g4
index f5116485..d40c6560 100644
--- a/src/main/antlr4/FIRRTL.g4
+++ b/src/main/antlr4/FIRRTL.g4
@@ -99,9 +99,9 @@ stmt
   | mdir 'mport' id '=' id '[' exp ']' exp info?
   | 'inst' id 'of' id info?
   | 'node' id '=' exp info?
-  | exp '<=' exp info?
-  | exp '<-' exp info?
-  | exp 'is' 'invalid' info?
+  | ref '<=' exp info?
+  | ref '<-' exp info?
+  | ref 'is' 'invalid' info?
   | when
   | 'stop(' exp exp intLit ')' stmtName? info?
   | 'printf(' exp exp StringLit ( exp)* ')' stmtName? info?
@@ -167,16 +167,22 @@ ruw
 exp
   : 'UInt' ('<' intLit '>')? '(' intLit ')'
   | 'SInt' ('<' intLit '>')? '(' intLit ')'
-  | id    // Ref
-  | exp '.' fieldId
-  | exp '.' DoubleLit // TODO Workaround for #470
-  | exp '[' intLit ']'
-  | exp '[' exp ']'
+  | ref
   | 'mux(' exp exp exp ')'
   | 'validif(' exp exp ')'
   | primop exp* intLit*  ')'
   ;
 
+ref
+  : id subref?
+  ;
+
+subref
+  : '.' fieldId subref?
+  | '.' DoubleLit subref? // TODO Workaround for #470
+  | '[' (intLit | exp) ']' subref?
+  ;
+
 id
   : Id
   | keywordAsId
diff --git a/src/main/scala/firrtl/Visitor.scala b/src/main/scala/firrtl/Visitor.scala
index ad2f5121..b2c4896f 100644
--- a/src/main/scala/firrtl/Visitor.scala
+++ b/src/main/scala/firrtl/Visitor.scala
@@ -6,6 +6,7 @@ import org.antlr.v4.runtime.ParserRuleContext
 import org.antlr.v4.runtime.tree.{AbstractParseTreeVisitor, ParseTreeVisitor, TerminalNode}
 import scala.collection.JavaConverters._
 import scala.collection.mutable
+import scala.annotation.tailrec
 import firrtl.antlr._
 import PrimOps._
 import FIRRTLParser._
@@ -441,9 +442,9 @@ class Visitor(infoMode: InfoMode) extends AbstractParseTreeVisitor[FirrtlNode] w
       // If we don't match on the first child, try the next one
       case _ =>
         ctx.getChild(1).getText match {
-          case "<=" => Connect(info, visitExp(ctx_exp(0)), visitExp(ctx_exp(1)))
-          case "<-" => PartialConnect(info, visitExp(ctx_exp(0)), visitExp(ctx_exp(1)))
-          case "is" => IsInvalid(info, visitExp(ctx_exp(0)))
+          case "<=" => Connect(info, visitRef(ctx.ref), visitExp(ctx_exp(0)))
+          case "<-" => PartialConnect(info, visitRef(ctx.ref), visitExp(ctx_exp(0)))
+          case "is" => IsInvalid(info, visitRef(ctx.ref))
           case "mport" =>
             CDefMPort(
               info,
@@ -457,32 +458,47 @@ class Visitor(infoMode: InfoMode) extends AbstractParseTreeVisitor[FirrtlNode] w
     }
   }
 
+  @tailrec private def visitSubRef(ctx: SubrefContext, inner: Expression): Expression = {
+    val ref = ctx.getChild(0).getText match {
+      case "." =>
+        if (ctx.fieldId != null) {
+          SubField(inner, ctx.fieldId.getText)
+        } else {
+          ctx.DoubleLit.getText.split('.') match {
+            case Array(a, b) if legalId(a) && legalId(b) => SubField(SubField(inner, a), b)
+            case _                                       => throw new ParserException(s"Illegal Expression at ${ctx.getText}")
+          }
+        }
+      case "[" =>
+        if (ctx.intLit != null) {
+          val lit = string2Int(ctx.intLit.getText)
+          SubIndex(inner, lit, UnknownType)
+        } else {
+          val idx = visitExp(ctx.exp)
+          SubAccess(inner, idx, UnknownType)
+        }
+    }
+    if (ctx.subref != null) {
+      visitSubRef(ctx.subref, ref)
+    } else {
+      ref
+    }
+  }
+
+  private def visitRef(ctx: RefContext): Expression = {
+    val ref = Reference(ctx.getChild(0).getText)
+    if (ctx.subref != null) {
+      visitSubRef(ctx.subref, ref)
+    } else {
+      ref
+    }
+  }
+
   private def visitExp(ctx: ExpContext): Expression = {
     val ctx_exp = ctx.exp.asScala
     ctx.getChild(0) match {
-      case _: IdContext => Reference(ctx.getText, UnknownType)
-      case _: ExpContext =>
-        ctx.getChild(1).getText match {
-          case "." =>
-            val expr1 = visitExp(ctx_exp(0))
-            // TODO Workaround for #470
-            if (ctx.fieldId == null) {
-              ctx.DoubleLit.getText.split('.') match {
-                case Array(a, b) if legalId(a) && legalId(b) =>
-                  val inner = new SubField(expr1, a, UnknownType)
-                  new SubField(inner, b, UnknownType)
-                case Array() => throw new ParserException(s"Illegal Expression at ${ctx.getText}")
-              }
-            } else {
-              new SubField(expr1, ctx.fieldId.getText, UnknownType)
-            }
-          case "[" =>
-            if (ctx.exp(1) == null)
-              new SubIndex(visitExp(ctx_exp(0)), string2Int(ctx.intLit(0).getText), UnknownType)
-            else
-              new SubAccess(visitExp(ctx_exp(0)), visitExp(ctx_exp(1)), UnknownType)
-        }
-      case _: PrimopContext =>
+      case ref: RefContext => visitRef(ref)
+      case _:   PrimopContext =>
         DoPrim(
           visitPrimop(ctx.primop),
           ctx_exp.map(visitExp).toSeq,