<commit_before>package com.boothen.jsonedit.model;

import org.antlr.v4.runtime.ANTLRInputStream;
import org.antlr.v4.runtime.CommonToken;
import org.antlr.v4.runtime.Lexer;
import org.antlr.v4.runtime.Token;
import org.eclipse.jface.text.BadLocationException;
import org.eclipse.jface.text.IDocument;
import org.eclipse.jface.text.TextAttribute;
import org.eclipse.jface.text.rules.IToken;
import org.eclipse.jface.text.rules.ITokenScanner;

import com.boothen.jsonedit.antlr.JSONLexer;
import com.boothen.jsonedit.core.JsonLog;
import com.boothen.jsonedit.core.preferences.TokenStyle;

/**
 * Use the ANTLR lexer to extract tokens from a document.
 */
public class AntlrTokenScanner implements ITokenScanner {

    private int offset;
    private Lexer lexer;
    private CommonToken previous;
    private CommonToken current;
    private TokenStyler tokenStyler;

    /**
     * Uses the given JSONLexer
     * @param lexer the lexer to use
     */
    public AntlrTokenScanner(Lexer lexer) {
        this.lexer = lexer;
        this.tokenStyler = new TokenStyler() {
            @Override
            public TextAttribute apply(TokenStyle style) {
                return new TextAttribute(null);
            }
        };
    }

    /**
     * @param lexer the lexer to use
     * @param mapping the mapping from token type to result
     */
    public AntlrTokenScanner(Lexer lexer, TokenStyler mapping) {
        this.lexer = lexer;
        this.tokenStyler = mapping;
    }

    @Override
    public void setRange(IDocument document, int offset, int length) {
        this.offset = offset;
        try {
            String text = document.get(offset, length);
            lexer.setInputStream(new ANTLRInputStream(text));
        } catch (BadLocationException e) {
            JsonLog.logError("Attempting to access a non-existing position", e);
        }
    }

    @Override
    public IToken nextToken() {
        Token token = lexer.nextToken();

        if (!(token instanceof CommonToken)) {
            // this doesn't seem to happen in practice
            // if it ever does: don't update current/previous, just bail out
            return org.eclipse.jface.text.rules.Token.UNDEFINED;
        }

        // update only if token is relevant (no WS)
        if (current == null || current.getType() != JSONLexer.WS) {
            previous = current;
        }

        current = (CommonToken) token;

        if (token.getType() == Token.EOF) {
            return org.eclipse.jface.text.rules.Token.EOF;
        }

        if (token.getType() == JSONLexer.WS) {
            return org.eclipse.jface.text.rules.Token.WHITESPACE;
        }

        int currentType = current.getType();
        int previousType = previous != null ? previous.getType() : Token.EOF;
        TokenStyle style = getStyle(currentType, previousType);
        TextAttribute data = tokenStyler.apply(style);

        return new org.eclipse.jface.text.rules.Token(data);
    }

    @Override
    public int getTokenOffset() {
        return offset + current.getStartIndex();
    }

    @Override
    public int getTokenLength() {
        return current.getStopIndex() - current.getStartIndex() + 1;
    }

    private TokenStyle getStyle(int currentTokenType, int previousTokenType) {
        switch (currentTokenType) {
        case JSONLexer.STRING:
            if (previousTokenType == JSONLexer.COLON) {
                return TokenStyle.TEXT;
            } else {
                return TokenStyle.KEY;
            }

        case JSONLexer.NUMBER:
            return TokenStyle.NUMBER;

        case JSONLexer.TRUE:
        case JSONLexer.FALSE:<commit_msg>Use two-step lookahead to corrently identify keys<commit_after>package com.boothen.jsonedit.model;

import org.antlr.v4.runtime.ANTLRInputStream;
import org.antlr.v4.runtime.CommonToken;
import org.antlr.v4.runtime.Lexer;
import org.antlr.v4.runtime.Token;
import org.eclipse.jface.text.BadLocationException;
import org.eclipse.jface.text.IDocument;
import org.eclipse.jface.text.TextAttribute;
import org.eclipse.jface.text.rules.IToken;
import org.eclipse.jface.text.rules.ITokenScanner;

import com.boothen.jsonedit.antlr.JSONLexer;
import com.boothen.jsonedit.core.JsonLog;
import com.boothen.jsonedit.core.preferences.TokenStyle;

/**
 * Use the ANTLR lexer to extract JSON tokens from a document.
 * It uses a 2-step lookahead to identify KEY text strings. It therefore does not require to
 * construct a syntax tree.
 */
public class AntlrTokenScanner implements ITokenScanner {

    private int offset;
    private Lexer lexer;
    private CommonToken current;
    private CommonToken next;
    private CommonToken afterNext;
    private TokenStyler tokenStyler;

    /**
     * Uses the given JSONLexer
     * @param lexer the lexer to use
     */
    public AntlrTokenScanner(Lexer lexer) {
        this.lexer = lexer;
        this.tokenStyler = new TokenStyler() {
            @Override
            public TextAttribute apply(TokenStyle style) {
                return new TextAttribute(null);
            }
        };
    }

    /**
     * @param lexer the lexer to use
     * @param mapping the mapping from token type to result
     */
    public AntlrTokenScanner(Lexer lexer, TokenStyler mapping) {
        this.lexer = lexer;
        this.tokenStyler = mapping;
    }

    @Override
    public void setRange(IDocument document, int offset, int length) {
        this.offset = offset;
        try {
            String text = document.get(offset, length);
            lexer.setInputStream(new ANTLRInputStream(text));
            next = (CommonToken) lexer.nextToken();
            afterNext = (CommonToken) lexer.nextToken();
        } catch (BadLocationException e) {
            JsonLog.logError("Attempting to access a non-existing position", e);
        }
    }

    @Override
    public IToken nextToken() {
        Token token = lexer.nextToken();

        if (!(token instanceof CommonToken)) {
            // this doesn't seem to happen in practice
            // if it ever does: don't update current/previous, just bail out
            return org.eclipse.jface.text.rules.Token.UNDEFINED;
        }

        current = next;
        next = afterNext;
        afterNext = (CommonToken) token;

        if (current.getType() == Token.EOF) {
            return org.eclipse.jface.text.rules.Token.EOF;
        }

        if (current.getType() == JSONLexer.WS) {
            return org.eclipse.jface.text.rules.Token.WHITESPACE;
        }

        int currentType = current.getType();
        int nextType = next.getType();
        int afterNextType = afterNext.getType();
        TokenStyle style = getStyle(currentType, nextType, afterNextType);
        TextAttribute data = tokenStyler.apply(style);

        return new org.eclipse.jface.text.rules.Token(data);
    }

    @Override
    public int getTokenOffset() {
        return offset + current.getStartIndex();
    }

    @Override
    public int getTokenLength() {
        return current.getStopIndex() - current.getStartIndex() + 1;
    }

    private TokenStyle getStyle(int currentType, int nextType, int afterNextType) {
        switch (currentType) {
        case JSONLexer.STRING:
            // There could be whitespace between current and the next (real) token
            // we therefore need to look two tokens ahead
            if (nextType == JSONLexer.COLON || (nextType == JSONLexer.WS && afterNextType == JSONLexer.COLON)) {
                return TokenStyle.KEY;
            } else {
                return TokenStyle.TEXT;
            }

        case JSONLexer.NUMBER:
            return TokenStyle.NUMBER;

        case JSONLexer.TRUE:
        case JSONLexer.FALSE:
{'fraction_non_alphanumeric': 0.09323132780082988, 'fraction_numerical': 0.0014263485477178422, 'mean_word_length': 2.3710664335664338, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3636457', 'n_tokens_mistral': 2276, 'n_tokens_neox': 2163, 'n_words': 554}