Retrieving rule element types by Context name instead of rule index

src/main/java/org/antlr/intellij/adaptor/parser/ANTLRParseTreeToPSIConverter.java

@@ -36,16 +36,22 @@ public class ANTLRParseTreeToPSIConverter implements ParseTreeListener {
 
 	protected final List<TokenIElementType> tokenElementTypes;
 	protected final List<RuleIElementType> ruleElementTypes;
+	private final IElementTypeMapper elementTypeMapper;
 
 	/** Map an error's start char index (usually start of a token) to the error object. */
 	protected Map<Integer, SyntaxError> tokenToErrorMap = new HashMap<>();
 
 	public ANTLRParseTreeToPSIConverter(Language language, Parser parser, PsiBuilder builder) {
+		this(language, parser, builder, new RuleIndexIElementTypeMapper(PSIElementTypeFactory.getRuleIElementTypes(language)));
+	}
+
+	public ANTLRParseTreeToPSIConverter(Language language, Parser parser, PsiBuilder builder, IElementTypeMapper elementTypeMapper) {
 		this.language = language;
 		this.builder = builder;
+		this.elementTypeMapper = elementTypeMapper;
 
 		this.tokenElementTypes = PSIElementTypeFactory.getTokenIElementTypes(language);
-		this.ruleElementTypes = PSIElementTypeFactory.getRuleIElementTypes(language);
+		this.ruleElementTypes = elementTypeMapper.getRuleElementTypes();
 
 		for (ANTLRErrorListener listener : parser.getErrorListeners()) {
 			if (listener instanceof SyntaxErrorListener) {
@@ -166,10 +172,10 @@ public void exitEveryRule(ParserRuleContext ctx) {
 			if (error != null) {
 				marker.error(error.getMessage());
 			} else {
-				marker.done(getRuleElementTypes().get(ctx.getRuleIndex()));
+				marker.done(elementTypeMapper.toIElementType(ctx));
 			}
 		} else {
-			marker.done(getRuleElementTypes().get(ctx.getRuleIndex()));
+			marker.done(elementTypeMapper.toIElementType(ctx));
 		}
 	}
 }

src/main/java/org/antlr/intellij/adaptor/parser/ClassNameIElementTypeMapper.java

@@ -0,0 +1,32 @@
+package org.antlr.intellij.adaptor.parser;
+
+import com.intellij.psi.tree.IElementType;
+import org.antlr.intellij.adaptor.lexer.RuleIElementType;
+import org.antlr.v4.runtime.ParserRuleContext;
+import org.apache.commons.lang.StringUtils;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class ClassNameIElementTypeMapper implements IElementTypeMapper {
+
+	private final Map<String, RuleIElementType> elementTypes;
+
+	public ClassNameIElementTypeMapper(List<RuleIElementType> elementTypes) {
+		this.elementTypes = elementTypes.stream()
+			.collect(Collectors.toMap(RuleIElementType::getDebugName, el -> el));
+	}
+
+	@Override
+	public IElementType toIElementType(ParserRuleContext ctx) {
+		String elementTypeName = StringUtils.removeEnd(ctx.getClass().getSimpleName(), "Context");
+		return elementTypes.get(elementTypeName);
+	}
+
+	@Override
+	public List<RuleIElementType> getRuleElementTypes() {
+		return new ArrayList<>(elementTypes.values());
+	}
+}

src/main/java/org/antlr/intellij/adaptor/parser/IElementTypeMapper.java

@@ -0,0 +1,16 @@
+package org.antlr.intellij.adaptor.parser;
+
+import com.intellij.psi.tree.IElementType;
+import org.antlr.intellij.adaptor.lexer.RuleIElementType;
+import org.antlr.v4.runtime.ParserRuleContext;
+
+import java.util.List;
+
+/**
+ * Maps a {@link org.antlr.v4.runtime.ParserRuleContext} to an {@link com.intellij.psi.tree.IElementType}.
+ */
+public interface IElementTypeMapper {
+	IElementType toIElementType(ParserRuleContext ctx);
+
+	List<RuleIElementType> getRuleElementTypes();
+}

src/main/java/org/antlr/intellij/adaptor/parser/RuleIndexIElementTypeMapper.java

@@ -0,0 +1,30 @@
+package org.antlr.intellij.adaptor.parser;
+
+import com.intellij.psi.tree.IElementType;
+import org.antlr.intellij.adaptor.lexer.RuleIElementType;
+import org.antlr.v4.runtime.ParserRuleContext;
+
+import java.util.List;
+
+/**
+ * Maps a {@link org.antlr.v4.runtime.ParserRuleContext} to an {@link com.intellij.psi.tree.IElementType} using
+ * a rule index.
+ */
+public class RuleIndexIElementTypeMapper implements IElementTypeMapper {
+
+	private final List<RuleIElementType> ruleElementTypes;
+
+	public RuleIndexIElementTypeMapper(List<RuleIElementType> elementTypes) {
+		this.ruleElementTypes = elementTypes;
+	}
+
+	@Override
+	public IElementType toIElementType(ParserRuleContext ctx) {
+		return ruleElementTypes.get(ctx.getRuleIndex());
+	}
+
+	@Override
+	public List<RuleIElementType> getRuleElementTypes() {
+		return ruleElementTypes;
+	}
+}