BlockEndState links to start state for each block in the ATN (constructed during deserialization)

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

@@ -236,6 +236,20 @@ public abstract class ATNSimulator {
 		}
 
 		for (ATNState state : atn.states) {
+			if (state instanceof BlockStartState) {
+				// we need to know the end state to set its start state
+				if (((BlockStartState)state).endState == null) {
+					throw new IllegalStateException();
+				}
+
+				// block end states can only be associated to a single block start state
+				if (((BlockStartState)state).endState.startState != null) {
+					throw new IllegalStateException();
+				}
+
+				((BlockStartState)state).endState.startState = (BlockStartState)state;
+			}
+
 			if (state instanceof PlusLoopbackState) {
 				PlusLoopbackState loopbackState = (PlusLoopbackState)state;
 				for (int i = 0; i < loopbackState.getNumberOfTransitions(); i++) {
@@ -310,6 +324,12 @@ public abstract class ATNSimulator {
 				}
 			}
 
+			if (state instanceof BlockEndState) {
+				if (((BlockEndState)state).startState == null) {
+					throw new IllegalStateException();
+				}
+			}
+
 			if (state instanceof DecisionState) {
 				DecisionState decisionState = (DecisionState)state;
 				if (decisionState.getNumberOfTransitions() > 1 && decisionState.decision < 0) {

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

@@ -31,4 +31,5 @@ package org.antlr.v4.runtime.atn;
 
 /** Terminal node of a simple (a|b|c) block */
 public class BlockEndState extends ATNState {
+	public BlockStartState startState;
 }