feat: allow retrieving rule element types by Context name instead of …

…rule index (#26)

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;
+ }
+}