From 2947fe6a2ad1c01f98920a8d42d668664719d23b Mon Sep 17 00:00:00 2001
From: Sam Harwell <sam@tunnelvisionlabs.com>
Date: Sat, 4 Aug 2012 22:22:28 -0500
Subject: [PATCH] Fix ATN optimizer leaving loop back state numbers in an
 inconsistent state

---
 .../src/org/antlr/v4/runtime/atn/ATNSimulator.java     | 10 +++++++++-
 .../src/org/antlr/v4/runtime/atn/LoopEndState.java     |  2 +-
 .../org/antlr/v4/runtime/atn/ParserATNSimulator.java   |  3 ++-
 tool/src/org/antlr/v4/automata/ATNSerializer.java      |  2 +-
 tool/src/org/antlr/v4/automata/ParserATNFactory.java   |  4 ++--
 5 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/runtime/Java/src/org/antlr/v4/runtime/atn/ATNSimulator.java b/runtime/Java/src/org/antlr/v4/runtime/atn/ATNSimulator.java
index e1f534f89..8cb3edfe7 100644
--- a/runtime/Java/src/org/antlr/v4/runtime/atn/ATNSimulator.java
+++ b/runtime/Java/src/org/antlr/v4/runtime/atn/ATNSimulator.java
@@ -32,6 +32,7 @@ package org.antlr.v4.runtime.atn;
 import org.antlr.v4.runtime.dfa.DFAState;
 import org.antlr.v4.runtime.misc.IntervalSet;
 import org.antlr.v4.runtime.misc.NotNull;
+import org.antlr.v4.runtime.misc.Pair;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -64,6 +65,7 @@ public abstract class ATNSimulator {
 		//
 		// STATES
 		//
+		List<Pair<LoopEndState, Integer>> loopBackStateNumbers = new ArrayList<Pair<LoopEndState, Integer>>();
 		int nstates = toInt(data[p++]);
 		for (int i=1; i<=nstates; i++) {
 			int stype = toInt(data[p++]);
@@ -75,11 +77,17 @@ public abstract class ATNSimulator {
 			ATNState s = stateFactory(stype, i);
 			s.ruleIndex = toInt(data[p++]);
 			if ( stype == ATNState.LOOP_END ) { // special case
-				((LoopEndState)s).loopBackStateNumber = toInt(data[p++]);
+				int loopBackStateNumber = toInt(data[p++]);
+				loopBackStateNumbers.add(new Pair<LoopEndState, Integer>((LoopEndState)s, loopBackStateNumber));
 			}
 			atn.addState(s);
 		}
 
+		// delay the assignment of loop back states until we know all the state instances have been initialized
+		for (Pair<LoopEndState, Integer> pair : loopBackStateNumbers) {
+			pair.a.loopBackState = atn.states.get(pair.b);
+		}
+
 		//
 		// RULES
 		//
diff --git a/runtime/Java/src/org/antlr/v4/runtime/atn/LoopEndState.java b/runtime/Java/src/org/antlr/v4/runtime/atn/LoopEndState.java
index cdde81ce5..a55330116 100644
--- a/runtime/Java/src/org/antlr/v4/runtime/atn/LoopEndState.java
+++ b/runtime/Java/src/org/antlr/v4/runtime/atn/LoopEndState.java
@@ -31,5 +31,5 @@ package org.antlr.v4.runtime.atn;
 
 /** Mark the end of a * or + loop */
 public class LoopEndState extends ATNState {
-	public int loopBackStateNumber;
+	public ATNState loopBackState;
 }
diff --git a/runtime/Java/src/org/antlr/v4/runtime/atn/ParserATNSimulator.java b/runtime/Java/src/org/antlr/v4/runtime/atn/ParserATNSimulator.java
index 4dca8ed11..0362bb648 100755
--- a/runtime/Java/src/org/antlr/v4/runtime/atn/ParserATNSimulator.java
+++ b/runtime/Java/src/org/antlr/v4/runtime/atn/ParserATNSimulator.java
@@ -987,7 +987,8 @@ public class ParserATNSimulator<Symbol extends Token> extends ATNSimulator {
 				if ( debug ) System.out.println("Loop end; pop, stack=" + config.context);
 				RuleContext p = config.context;
 				LoopEndState end = (LoopEndState) config.state;
-				while ( !p.isEmpty() && p.invokingState == end.loopBackStateNumber ) {
+				int loopBackStateNumber = end.loopBackState.stateNumber;
+				while ( !p.isEmpty() && p.invokingState == loopBackStateNumber ) {
 					p = config.context = config.context.parent; // "pop"
 				}
 			}
diff --git a/tool/src/org/antlr/v4/automata/ATNSerializer.java b/tool/src/org/antlr/v4/automata/ATNSerializer.java
index bb3fa6653..e83855194 100644
--- a/tool/src/org/antlr/v4/automata/ATNSerializer.java
+++ b/tool/src/org/antlr/v4/automata/ATNSerializer.java
@@ -89,7 +89,7 @@ public class ATNSerializer {
 			}
 			data.add(s.getStateType());
 			data.add(s.ruleIndex);
-			if ( s.getStateType() == ATNState.LOOP_END ) data.add(((LoopEndState)s).loopBackStateNumber);
+			if ( s.getStateType() == ATNState.LOOP_END ) data.add(((LoopEndState)s).loopBackState.stateNumber);
 			nedges += s.getNumberOfTransitions();
 			for (int i=0; i<s.getNumberOfTransitions(); i++) {
 				Transition t = s.transition(i);
diff --git a/tool/src/org/antlr/v4/automata/ParserATNFactory.java b/tool/src/org/antlr/v4/automata/ParserATNFactory.java
index b05ffc695..376846e6f 100644
--- a/tool/src/org/antlr/v4/automata/ParserATNFactory.java
+++ b/tool/src/org/antlr/v4/automata/ParserATNFactory.java
@@ -425,7 +425,7 @@ public class ParserATNFactory implements ATNFactory {
 		atn.defineDecisionState(loop);
 		LoopEndState end = newState(LoopEndState.class, plusAST);
 		blkStart.loopBackState = loop;
-		end.loopBackStateNumber = loop.stateNumber;
+		end.loopBackState = loop;
 
 		plusAST.atnState = blkStart;
 		epsilon(blkEnd, loop);		// blk can see loop back
@@ -468,7 +468,7 @@ public class ParserATNFactory implements ATNFactory {
 		LoopEndState end = newState(LoopEndState.class, starAST);
 		StarLoopbackState loop = newState(StarLoopbackState.class, starAST);
 		entry.loopBackState = loop;
-		end.loopBackStateNumber = loop.stateNumber;
+		end.loopBackState = loop;
 
 		BlockAST blkAST = (BlockAST)starAST.getChild(0);
 		entry.isGreedy = isGreedy(blkAST);