refactor

runtime-testsuite/test/org/antlr/v4/runtime/TestSuggestionHelper.java

@@ -56,7 +56,7 @@ void matchesExpectation(Path expectations) throws Exception {
 		JavaParser parser = new JavaParser(stream);
 		parser.setStateListener(new TokenStateRecorder(stream));
 		RuleContext context = parser.compilationUnit();
-		Pair<RuleContext, IntervalSet> contextsAndIntervals = parser.getExpectedTokens(context, expectation.line, expectation.column);
+		Pair<RuleContext, IntervalSet> contextsAndIntervals = parser.getExpectedTokensAt(context, expectation.line, expectation.column);
 		Set<String> actual = contextsAndIntervals.b.toSet().stream()
 			.map(t -> parser.getVocabulary().getDisplayName(t))
 			.map(s -> s.startsWith("'") ? s.substring(1, s.length() - 1) : s)

runtime/Java/src/org/antlr/v4/runtime/Parser.java

@@ -808,12 +808,12 @@ public IntervalSet getExpectedTokens() {
 		return getATN().getExpectedTokens(getState(), getContext());
 	}
 
-	public Pair<RuleContext, IntervalSet> getExpectedTokens(RuleContext startRuleContext, int line, int column) {
-		Token lastToken = this._input.locateTokenAtOrBefore(line, column);
+	public Pair<RuleContext, IntervalSet> getExpectedTokensAt(RuleContext startRuleContext, int line, int column) {
+		Token lastToken = this._input.lastTokenAt(line, column);
 		if(!(lastToken instanceof TokenWithStates)) {
 			throw new IllegalStateException("Cannot getExpectedTokens. Configure Lexer with CommonTokenWithStateFactory and parser with TokenStateRecorder");
 		}
-		RuleContext context = locateContextWithToken(startRuleContext, lastToken);
+		RuleContext context = startRuleContext.leafContextWithToken(lastToken);
 		TokenWithStates token = (TokenWithStates)lastToken;
 		int stateNumber = token.getFollowState();
 		if(stateNumber == ATNState.INVALID_STATE_NUMBER)
@@ -824,28 +824,6 @@ public Pair<RuleContext, IntervalSet> getExpectedTokens(RuleContext startRuleCon
 		return new Pair<>(context, set);
 	}
 
-	private RuleContext locateContextWithToken(ParseTree tree, Token token) {
-		for(int i=0; i<tree.getChildCount(); i++) {
-			ParseTree child = tree.getChild(i);
-			Interval tokens = child.getSourceInterval();
-			// skip empty interval
-			if(tokens.b < tokens.a)
-				continue;
-			// have we gone past token ?
-			if(tokens.a > token.getTokenIndex())
-				break;
-			if (tokens.b >= token.getTokenIndex()) {
-				RuleContext context = locateContextWithToken(child, token);
-				if(context!=null)
-					return context;
-			}
-		}
-		if(tree instanceof RuleNode)
-			return ((RuleNode) tree).getRuleContext();
-		else
-			return null;
-	}
-
 	public IntervalSet getExpectedTokensWithinCurrentRule() {
         ATN atn = getInterpreter().atn;
         ATNState s = atn.states.get(getState());

runtime/Java/src/org/antlr/v4/runtime/TokenStream.java

@@ -138,7 +138,7 @@ public interface TokenStream extends IntStream {
 	 */
 	public String getText(Token start, Token stop);
 
-    public default Token locateTokenAtOrBefore(int line, int column) {
+    public default Token lastTokenAt(int line, int column) {
 		Token last = get(0);
 		for(int i=1; i<index(); i++) {
 			Token token = get(i);

runtime/Java/src/org/antlr/v4/runtime/tree/SyntaxTree.java

@@ -6,6 +6,8 @@
 
 package org.antlr.v4.runtime.tree;
 
+import org.antlr.v4.runtime.RuleContext;
+import org.antlr.v4.runtime.Token;
 import org.antlr.v4.runtime.TokenStream;
 import org.antlr.v4.runtime.misc.Interval;
 
@@ -31,4 +33,27 @@ public interface SyntaxTree extends Tree {
 	 * EOF is unspecified.</p>
 	 */
 	Interval getSourceInterval();
+
+	default RuleContext leafContextWithToken(Token token) {
+		for(int i=0; i<getChildCount(); i++) {
+			SyntaxTree child = (SyntaxTree)getChild(i);
+			Interval tokens = child.getSourceInterval();
+			// skip empty interval
+			if(tokens.b < tokens.a)
+				continue;
+			// have we gone past token ?
+			if(tokens.a > token.getTokenIndex())
+				break;
+			if (tokens.b >= token.getTokenIndex()) {
+				RuleContext context = child.leafContextWithToken(token);
+				if(context!=null)
+					return context;
+			}
+		}
+		if(this instanceof RuleNode)
+			return ((RuleNode) this).getRuleContext();
+		else
+			return null;
+
+	}
 }