Fix error highlighting for missing tokens

antlr#2
ksprojects/protobuf-jetbrains-plugin#25

src/main/java/org/antlr/jetbrains/adapter/parser/ANTLRParseTreeToPSIConverter.java

@@ -6,6 +6,7 @@
 import org.antlr.jetbrains.adapter.lexer.PSIElementTypeFactory;
 import org.antlr.jetbrains.adapter.lexer.RuleIElementType;
 import org.antlr.jetbrains.adapter.lexer.TokenIElementType;
+import org.antlr.v4.runtime.RecognitionException;
 import org.antlr.v4.runtime.ANTLRErrorListener;
 import org.antlr.v4.runtime.Parser;
 import org.antlr.v4.runtime.ParserRuleContext;
@@ -30,7 +31,7 @@
 public class ANTLRParseTreeToPSIConverter implements ParseTreeListener {
 	protected final Language language;
 	protected final PsiBuilder builder;
-	protected List<SyntaxError> syntaxErrors;
+	protected Map<RecognitionException, SyntaxError> syntaxErrors;
 	protected final Deque<PsiBuilder.Marker> markers = new ArrayDeque<PsiBuilder.Marker>();
 
 	protected final List<TokenIElementType> tokenElementTypes;
@@ -48,8 +49,8 @@ public ANTLRParseTreeToPSIConverter(Language language, Parser parser, PsiBuilder
 
 		for (ANTLRErrorListener listener : parser.getErrorListeners()) {
 			if (listener instanceof SyntaxErrorListener) {
-				syntaxErrors = ((SyntaxErrorListener)listener).getSyntaxErrors();
-				for (SyntaxError error : syntaxErrors) {
+				syntaxErrors = ((SyntaxErrorListener)listener).getErrorMap();
+				for (SyntaxError error : syntaxErrors.values()) {
 					// record first error per token
 					int StartIndex = error.getOffendingSymbol().getStartIndex();
 					if ( !tokenToErrorMap.containsKey(StartIndex) ) {
@@ -115,6 +116,7 @@ public void visitTerminal(TerminalNode node) {
 	 *  prediction started (which we use to find syntax errors). So,
 	 *  SyntaxError objects return start not offending token in this case.
 	 */
+	@Override
 	public void visitErrorNode(ErrorNode node) {
 		ProgressIndicatorProvider.checkCanceled();
 
@@ -159,6 +161,16 @@ public void enterEveryRule(ParserRuleContext ctx) {
 	public void exitEveryRule(ParserRuleContext ctx) {
 		ProgressIndicatorProvider.checkCanceled();
 		PsiBuilder.Marker marker = markers.pop();
-		marker.done(getRuleElementTypes().get(ctx.getRuleIndex()));
+		if (ctx.exception != null) {
+			SyntaxError error = syntaxErrors.get(ctx.exception);
+			if (error != null) {
+				marker.error(error.getMessage());
+			} else {
+				// should not happen
+				marker.error("syntax error");
+			}
+		} else {
+			marker.done(getRuleElementTypes().get(ctx.getRuleIndex()));
+		}
 	}
 }

src/main/java/org/antlr/jetbrains/adapter/parser/SyntaxErrorListener.java

@@ -6,20 +6,23 @@
 import org.antlr.v4.runtime.Token;
 import org.antlr.v4.runtime.misc.Utils;
 
-import java.util.ArrayList;
-import java.util.List;
+import java.util.*;
 
 /** Traps errors from parsing language of plugin. E.g., for a Java plugin,
  *  this would catch errors when people type invalid Java code into .java file.
  *  This swallows the errors as the PSI tree has error nodes.
  */
 public class SyntaxErrorListener extends BaseErrorListener {
-	private final List<SyntaxError> syntaxErrors = new ArrayList<SyntaxError>();
+	private final Map<RecognitionException, SyntaxError> syntaxErrors = new HashMap<RecognitionException, SyntaxError>();
 
 	public SyntaxErrorListener() {
 	}
 
 	public List<SyntaxError> getSyntaxErrors() {
+		return new ArrayList<SyntaxError>(syntaxErrors.values());
+	}
+
+	public Map<RecognitionException, SyntaxError> getErrorMap() {
 		return syntaxErrors;
 	}
 
@@ -29,11 +32,13 @@ public void syntaxError(Recognizer<?, ?> recognizer,
 							int line, int charPositionInLine,
 							String msg, RecognitionException e)
 	{
-		syntaxErrors.add(new SyntaxError(recognizer, (Token)offendingSymbol, line, charPositionInLine, msg, e));
+		syntaxErrors.put(e, new SyntaxError(recognizer, (Token)offendingSymbol, line, charPositionInLine, msg, e));
 	}
 
+
+
 	@Override
 	public String toString() {
-		return Utils.join(syntaxErrors.iterator(), "\n");
+		return Utils.join(syntaxErrors.values().iterator(), "\n");
 	}
 }