rm'd epsilon transition in ATN after rule refs.

[git-p4: depot-paths = "//depot/code/antlr4/main/": change = 9518]

runtime/Java/src/org/antlr/v4/runtime/atn/RuleTransition.java

@@ -38,7 +38,7 @@ public class RuleTransition extends Transition {
 
 	/** What node to begin computations following ref to rule */
 	@NotNull
-	public final ATNState followState;
+	public ATNState followState;
 
 	public RuleTransition(@NotNull RuleStartState ruleStart,
 						  int ruleIndex,

tool/playground/T.g

@@ -1,8 +1,4 @@
 grammar T;
-s : b ';' |  b '.' ;
+s : b c ;
-b : a ;
+b : A ;
-a  : {false}? ID {System.out.println("alt 1");}
+c : C ;
-  | {true}? ID {System.out.println("alt 2");}
-  ;ID : 'a'..'z'+ ;
-INT : '0'..'9'+;
-WS : (' '|'\n') {skip();} ;

tool/src/org/antlr/v4/automata/ParserATNFactory.java

@@ -58,11 +58,26 @@ import java.util.List;
  *  No side-effects. It builds an ATN object and returns it.
  */
 public class ParserATNFactory implements ATNFactory {
-	class TailEpsilonRemover extends ATNVisitor {
+    /** Add follow links from rule stop states to every following state for rule invocations.
-		@Override
+     *  Must do after construction since we optimize away some epsilon transitions.
-		public void visitState(ATNState p) {
+     */
-			if ( p.getClass() == ATNState.class && p.getNumberOfTransitions()==1 ) {
+    class FollowLinkAdder extends ATNVisitor {
-				ATNState q = p.transition(0).target;
+        @Override
+        public void visitState(ATNState p) {
+            if ( p.getClass() == ATNState.class && p.getNumberOfTransitions()==1 &&
+                 p.transition(0) instanceof RuleTransition )
+            {
+                RuleTransition rt = (RuleTransition) p.transition(0);
+                addFollowLink(rt.ruleIndex, rt.followState);
+            }
+        }
+    }
+
+    class TailEpsilonRemover extends ATNVisitor {
+        @Override
+        public void visitState(ATNState p) {
+            if ( p.getClass() == ATNState.class && p.getNumberOfTransitions()==1 ) {
+                ATNState q = p.transition(0).target;
 				if ( q.getClass() == ATNState.class ) {
 					// we have p-x->q for x in {rule, action, pred, token, ...}
 					// if edge out of q is single epsilon to block end
@@ -96,11 +111,12 @@ public class ParserATNFactory implements ATNFactory {
 	public ATN createATN() {
 		_createATN(g.rules.values());
 		atn.maxTokenType = g.getMaxTokenType();
+        addRuleFollowLinks();
 		addEOFTransitionToStartRules();
 		return atn;
 	}
 
-	public void _createATN(Collection<Rule> rules) {
+    public void _createATN(Collection<Rule> rules) {
 		createRuleStartAndStopATNStates();
 
 		GrammarASTAdaptor adaptor = new GrammarASTAdaptor();
@@ -205,8 +221,8 @@ public class ParserATNFactory implements ATNFactory {
 	 */
 	public Handle ruleRef(GrammarAST node) {
 		Handle h = _ruleRef(node);
-		Rule r = g.getRule(node.getText());
+//		Rule r = g.getRule(node.getText());
-		addFollowLink(r, h.right);
+//		addFollowLink(r, h.right);
 		return h;
 	}
 
@@ -222,9 +238,9 @@ public class ParserATNFactory implements ATNFactory {
 		return new Handle(left, right);
 	}
 
-	public void addFollowLink(Rule r, ATNState right) {
+	public void addFollowLink(int ruleIndex, ATNState right) {
 		// add follow edge from end of invoked rule
-		RuleStopState stop = atn.ruleToStopState[r.index];
+		RuleStopState stop = atn.ruleToStopState[ruleIndex];
 		epsilon(stop, right);
 	}
 
@@ -349,13 +365,16 @@ public class ParserATNFactory implements ATNFactory {
 			Handle el = els.get(i);
 			// if el is of form o-x->o for x in {rule, action, pred, token, ...}
 			// and not last in alt
-			if ( el.left.getClass() == ATNState.class &&
+            Transition tr = null;
+            if ( el.left.getNumberOfTransitions()==1 ) tr = el.left.transition(0);
+            boolean isRuleTrans = tr instanceof RuleTransition;
+            if ( el.left.getClass() == ATNState.class &&
 				el.right.getClass() == ATNState.class &&
-				el.left.getNumberOfTransitions()==1 &&
+				tr!=null && (isRuleTrans || tr.target == el.right) )
-				el.left.transition(0).target == el.right)
 			{
 				// we can avoid epsilon edge to next el
-				el.left.transition(0).target = els.get(i+1).left;
+				if ( isRuleTrans ) ((RuleTransition)tr).followState = els.get(i+1).left;
+                else tr.target = els.get(i+1).left;
 				atn.removeState(el.right); // we skipped over this state
 			}
 			else { // need epsilon if previous block's right end node is complicated
@@ -501,6 +520,14 @@ public class ParserATNFactory implements ATNFactory {
 		}
 	}
 
+    public void addRuleFollowLinks() {
+        FollowLinkAdder flinker = new FollowLinkAdder();
+        for (Rule r : g.rules.values()) {
+        	ATNState start = atn.ruleToStartState[r.index];
+            flinker.visit(start);
+        }
+    }
+
 	/** add an EOF transition to any rule end ATNState that points to nothing
      *  (i.e., for all those rules not invoked by another rule).  These
      *  are start symbols then.

tool/test/org/antlr/v4/test/TestATNConstruction.java

@@ -184,17 +184,16 @@ public class TestATNConstruction extends BaseTest {
 			"b : B ;");
 		String expecting =
 			"RuleStart_a_0->s4\n" +
-			"s4-b->RuleStart_b_2\n" +
+            "s4-b->RuleStart_b_2\n" +
-			"s5->s6\n" +
+            "s6-A->s7\n" +
-			"s6-A->s7\n" +
+            "s7->RuleStop_a_1\n" +
-			"s7->RuleStop_a_1\n" +
+            "RuleStop_a_1-EOF->s10\n";
-			"RuleStop_a_1-EOF->s10\n";
 		checkRuleATN(g, "a", expecting);
 		expecting =
 			"RuleStart_b_2->s8\n" +
 			"s8-B->s9\n" +
 			"s9->RuleStop_b_3\n" +
-			"RuleStop_b_3->s5\n";
+			"RuleStop_b_3->s6\n";
 		checkRuleATN(g, "b", expecting);
 	}
 
@@ -206,10 +205,10 @@ public class TestATNConstruction extends BaseTest {
 			"c : b C;");
 		String expecting =
 			"RuleStart_b_2->s10\n" +
-			"s10-B->s11\n" +
+            "s10-B->s11\n" +
-			"s11->RuleStop_b_3\n" +
+            "s11->RuleStop_b_3\n" +
-			"RuleStop_b_3->s7\n" +
+            "RuleStop_b_3->s8\n" +
-			"RuleStop_b_3->s13\n";
+            "RuleStop_b_3->s14\n";
 		checkRuleATN(g, "b", expecting);
 	}