From 0981a397b53860a7f40999952f6b2c7f9df2615e Mon Sep 17 00:00:00 2001
From: Sam Harwell <sam.harwell@rackspace.com>
Date: Wed, 20 Nov 2013 09:09:26 -0600
Subject: [PATCH] Add option to generate rule bypass transitions during ATN
 deserialization

---
 .../src/org/antlr/v4/runtime/atn/ATN.java     |  6 +-
 .../atn/ATNDeserializationOptions.java        | 12 +++
 .../antlr/v4/runtime/atn/ATNDeserializer.java | 90 +++++++++++++++++++
 3 files changed, 106 insertions(+), 2 deletions(-)

diff --git a/runtime/Java/src/org/antlr/v4/runtime/atn/ATN.java b/runtime/Java/src/org/antlr/v4/runtime/atn/ATN.java
index db852a081..dd67ba086 100644
--- a/runtime/Java/src/org/antlr/v4/runtime/atn/ATN.java
+++ b/runtime/Java/src/org/antlr/v4/runtime/atn/ATN.java
@@ -82,8 +82,10 @@ public class ATN {
 
 	/**
 	 * For lexer ATNs, this maps the rule index to the resulting token type.
-	 * <p/>
-	 * This is {@code null} for parser ATNs.
+	 * For parser ATNs, this maps the rule index to the generated bypass token
+	 * type if the
+	 * {@link ATNDeserializationOptions#isGenerateRuleBypassTransitions}
+	 * deserialization option was specified; otherwise, this is {@code null}.
 	 */
 	public int[] ruleToTokenType;
 
diff --git a/runtime/Java/src/org/antlr/v4/runtime/atn/ATNDeserializationOptions.java b/runtime/Java/src/org/antlr/v4/runtime/atn/ATNDeserializationOptions.java
index 123199b65..a1fd40d04 100644
--- a/runtime/Java/src/org/antlr/v4/runtime/atn/ATNDeserializationOptions.java
+++ b/runtime/Java/src/org/antlr/v4/runtime/atn/ATNDeserializationOptions.java
@@ -45,13 +45,16 @@ public class ATNDeserializationOptions {
 
 	private boolean readOnly;
 	private boolean verifyATN;
+	private boolean generateRuleBypassTransitions;
 
 	public ATNDeserializationOptions() {
 		this.verifyATN = true;
+		this.generateRuleBypassTransitions = false;
 	}
 
 	public ATNDeserializationOptions(ATNDeserializationOptions options) {
 		this.verifyATN = options.verifyATN;
+		this.generateRuleBypassTransitions = options.generateRuleBypassTransitions;
 	}
 
 	@NotNull
@@ -76,6 +79,15 @@ public class ATNDeserializationOptions {
 		this.verifyATN = verifyATN;
 	}
 
+	public final boolean isGenerateRuleBypassTransitions() {
+		return generateRuleBypassTransitions;
+	}
+
+	public final void setGenerateRuleBypassTransitions(boolean generateRuleBypassTransitions) {
+		throwIfReadOnly();
+		this.generateRuleBypassTransitions = generateRuleBypassTransitions;
+	}
+
 	protected void throwIfReadOnly() {
 		if (isReadOnly()) {
 			throw new IllegalStateException("The object is read only.");
diff --git a/runtime/Java/src/org/antlr/v4/runtime/atn/ATNDeserializer.java b/runtime/Java/src/org/antlr/v4/runtime/atn/ATNDeserializer.java
index 0814974a0..a3fd7606c 100644
--- a/runtime/Java/src/org/antlr/v4/runtime/atn/ATNDeserializer.java
+++ b/runtime/Java/src/org/antlr/v4/runtime/atn/ATNDeserializer.java
@@ -362,6 +362,96 @@ public class ATNDeserializer {
 			verifyATN(atn);
 		}
 
+		if (deserializationOptions.isGenerateRuleBypassTransitions() && atn.grammarType == ATNType.PARSER) {
+			atn.ruleToTokenType = new int[atn.ruleToStartState.length];
+			for (int i = 0; i < atn.ruleToStartState.length; i++) {
+				atn.ruleToTokenType[i] = atn.maxTokenType + i + 1;
+			}
+
+			for (int i = 0; i < atn.ruleToStartState.length; i++) {
+				BasicBlockStartState bypassStart = new BasicBlockStartState();
+				bypassStart.ruleIndex = i;
+				atn.addState(bypassStart);
+
+				BlockEndState bypassStop = new BlockEndState();
+				bypassStop.ruleIndex = i;
+				atn.addState(bypassStop);
+
+				bypassStart.endState = bypassStop;
+				atn.defineDecisionState(bypassStart);
+
+				bypassStop.startState = bypassStart;
+
+				ATNState endState;
+				Transition excludeTransition = null;
+				if (atn.ruleToStartState[i].isPrecedenceRule) {
+					// wrap from the beginning of the rule to the StarLoopEntryState
+					endState = null;
+					for (ATNState state : atn.states) {
+						if (state.ruleIndex != i) {
+							continue;
+						}
+
+						if (!(state instanceof StarLoopEntryState)) {
+							continue;
+						}
+
+						ATNState maybeLoopEndState = state.transition(state.getNumberOfTransitions() - 1).target;
+						if (!(maybeLoopEndState instanceof LoopEndState)) {
+							continue;
+						}
+
+						if (maybeLoopEndState.epsilonOnlyTransitions && maybeLoopEndState.transition(0).target instanceof RuleStopState) {
+							endState = state;
+							break;
+						}
+					}
+
+					if (endState == null) {
+						throw new UnsupportedOperationException("Couldn't identify final state of the precedence rule prefix section.");
+					}
+
+					excludeTransition = ((StarLoopEntryState)endState).loopBackState.transition(0);
+				}
+				else {
+					endState = atn.ruleToStopState[i];
+				}
+
+				// all non-excluded transitions that currently target end state need to target blockEnd instead
+				for (ATNState state : atn.states) {
+					for (Transition transition : state.transitions) {
+						if (transition == excludeTransition) {
+							continue;
+						}
+
+						if (transition.target == endState) {
+							transition.target = bypassStop;
+						}
+					}
+				}
+
+				// all transitions leaving the rule start state need to leave blockStart instead
+				while (atn.ruleToStartState[i].getNumberOfTransitions() > 0) {
+					Transition transition = atn.ruleToStartState[i].removeTransition(atn.ruleToStartState[i].getNumberOfTransitions() - 1);
+					bypassStart.addTransition(transition);
+				}
+
+				// link the new states
+				atn.ruleToStartState[i].addTransition(new EpsilonTransition(bypassStart));
+				bypassStop.addTransition(new EpsilonTransition(endState));
+
+				ATNState matchState = new BasicState();
+				atn.addState(matchState);
+				matchState.addTransition(new AtomTransition(bypassStop, atn.ruleToTokenType[i]));
+				bypassStart.addTransition(new EpsilonTransition(matchState));
+			}
+
+			if (deserializationOptions.isVerifyATN()) {
+				// reverify after modification
+				verifyATN(atn);
+			}
+		}
+
 		return atn;
 	}