Retrieve token list from ParserRuleContext

Question

Main question

Is there an easy way to get a list of tokens (ideally in a form of a TokenStream) from the parser rule class ParserRuleContext?

Related answers

In an answer for a question Traversal of tokens using ParserRuleContext in listener - ANTLR4 this solution appeared:

ParserRuleContext pctx = ctx.getParent();
List<TerminalNode> nodes = pctx.getTokens(pctx.getStart(), pctx.getStop());

But there is no method with signature ParserRuleContext::getTokens(Token, Token) in ANTLRv4.

My solution

I thought about retriving a list of tokens from TokenStream by using TokenStream:get(index: int) method, where index value will be set to a range of indicies of given ParserRuleContext start/stop tokens.

Additional question

Is there a way to get a subset of tokens from TokenStream in a form of another TokenStream?

Answer 1

So, I overlooked some classes and their inferfaces in the ANTLRv4 API.

Answer to main question

Proposed above solution is correct. Also BufferedTokenStream and CommonTokenStream classes have method public List<Token> getTokens(int start, int stop) which allows to retrive list of tokens from a given range (especially from a range between start token and stop token of ParserRuleContext class)

Answer to additional question

You can utilize ListTokenSource class which implements TokenSource interface. Then you can create CommonTokenStream class passing the ListTokenSource.

Code example of ParserRuleRewriter

I encapsulate above ideas into small code example featuring ParserRuleRewriter - a TokenStreamRewriter that rewrites only given parser rule. In the code tokenStream parameter is a token stream of a full program.

import org.antlr.v4.runtime.*;
import java.util.List;

public class ParserRuleRewriter {
    private TokenStreamRewriter rewriter;

    public ParserRuleRewriter(ParserRuleContext parserRule, CommonTokenStream tokenStream) {
        Token start = parserRule.getStart();
        Token stop = parserRule.getStop();
        List<Token> ruleTokens = tokenStream.getTokens(start.getTokenIndex(), stop.getTokenIndex());
        ListTokenSource tokenSource = new ListTokenSource(ruleTokens);
        CommonTokenStream commonTokenStream = new CommonTokenStream(tokenSource);
        commonTokenStream.fill();
        rewriter = new TokenStreamRewriter(commonTokenStream);
    }

    public void replace(Token token, ParserRuleContext rule) {
        rewriter.replace(token, rule.getText());
    }

    @Override
    public String toString() {
        return rewriter.getText();
    }
}