From bbe77782db1f8e6dbbf4204293ffdce620344365 Mon Sep 17 00:00:00 2001
From: Sam Harwell <sam@tunnelvisionlabs.com>
Date: Mon, 15 Oct 2012 00:56:25 -0500
Subject: [PATCH] Fix state checking in new non-greedy handling

---
 .../antlr/v4/runtime/atn/LexerATNConfig.java  |  2 +-
 .../v4/runtime/atn/LexerATNSimulator.java     | 35 ++++++++++++++++---
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/runtime/Java/src/org/antlr/v4/runtime/atn/LexerATNConfig.java b/runtime/Java/src/org/antlr/v4/runtime/atn/LexerATNConfig.java
index 5cf5a1dbc..d3f857d85 100644
--- a/runtime/Java/src/org/antlr/v4/runtime/atn/LexerATNConfig.java
+++ b/runtime/Java/src/org/antlr/v4/runtime/atn/LexerATNConfig.java
@@ -76,7 +76,7 @@ public class LexerATNConfig extends ATNConfig {
 	}
 
 	public LexerATNConfig exitNonGreedyBlock() {
-		if (!isGreedy()) {
+		if (isGreedy()) {
 			return this;
 		}
 
diff --git a/runtime/Java/src/org/antlr/v4/runtime/atn/LexerATNSimulator.java b/runtime/Java/src/org/antlr/v4/runtime/atn/LexerATNSimulator.java
index 3d1439cf7..95d39f7f8 100644
--- a/runtime/Java/src/org/antlr/v4/runtime/atn/LexerATNSimulator.java
+++ b/runtime/Java/src/org/antlr/v4/runtime/atn/LexerATNSimulator.java
@@ -417,6 +417,7 @@ public class LexerATNSimulator extends ATNSimulator {
 			}
 
 			if (!((LexerATNConfig)config).isGreedy()) {
+				assert !(config.state instanceof RuleStopState);
 				nonGreedyAlts.set(config.alt);
 			}
 		}
@@ -602,13 +603,13 @@ public class LexerATNSimulator extends ATNSimulator {
 		}
 
 		ATNState p = config.state;
+		boolean nonGreedy = p instanceof DecisionState && ((DecisionState)p).nonGreedy
+			|| p instanceof PlusBlockStartState && ((PlusBlockStartState)p).loopBackState.nonGreedy;
 		for (int i=0; i<p.getNumberOfTransitions(); i++) {
 			Transition t = p.transition(i);
 			LexerATNConfig c = getEpsilonTarget(config, t, configs);
 			if ( c!=null ) {
-				final int NON_GREEDY_ENTER_ALT = 2;
-				if (i == NON_GREEDY_ENTER_ALT - 1 && ((DecisionState)p).nonGreedy) {
-					assert p.getNumberOfTransitions() == 2;
+				if (nonGreedy) {
 					c = c.enterNonGreedyBlock();
 				}
 
@@ -624,8 +625,34 @@ public class LexerATNSimulator extends ATNSimulator {
 									  @NotNull ATNConfigSet configs)
 	{
 		ATNState p = config.state;
-		if (p.isNonGreedyExitState()) {
+		switch (p.getStateType()) {
+		case ATNState.PLUS_LOOP_BACK:
+			if (t.target instanceof PlusBlockStartState && ((PlusBlockStartState)t.target).nonGreedy) {
+				config = config.exitNonGreedyBlock();
+			}
+
+			break;
+
+		case ATNState.STAR_LOOP_BACK:
 			config = config.exitNonGreedyBlock();
+			break;
+
+		case ATNState.LOOP_END:
+			ATNState loopEntry = ((LoopEndState)p).loopBackState.transition(0).target;
+			if (loopEntry instanceof StarLoopEntryState && ((StarLoopEntryState)loopEntry).nonGreedy) {
+				config = config.exitNonGreedyBlock();
+			}
+
+			break;
+
+		case ATNState.BLOCK_END:
+			if (p.isNonGreedyExitState()) {
+				config = config.exitNonGreedyBlock();
+			}
+			break;
+
+		default:
+			break;
 		}
 
 		LexerATNConfig c = null;