ScotchLayoutLexer.java

package scotch.compiler.parser;

import static java.util.Arrays.asList;
import static org.antlr.v4.runtime.Token.EOF;
import static scotch.compiler.parser.ScotchLayoutLexer.State.ACCEPT;
import static scotch.compiler.parser.ScotchLayoutLexer.State.SCAN_DEFAULT;
import static scotch.compiler.parser.ScotchLayoutLexer.State.SCAN_DISABLED;
import static scotch.compiler.parser.ScotchLayoutLexer.State.SCAN_LAYOUT;
import static scotch.compiler.parser.ScotchLayoutLexer.State.SCAN_LET;
import static scotch.compiler.parser.ScotchLexer.CLOSE_CURLY;
import static scotch.compiler.parser.ScotchLexer.CLOSE_SQUARE;
import static scotch.compiler.parser.ScotchLexer.DO;
import static scotch.compiler.parser.ScotchLexer.IN;
import static scotch.compiler.parser.ScotchLexer.LET;
import static scotch.compiler.parser.ScotchLexer.NEWLINE;
import static scotch.compiler.parser.ScotchLexer.ON;
import static scotch.compiler.parser.ScotchLexer.OPEN_CURLY;
import static scotch.compiler.parser.ScotchLexer.OPEN_SQUARE;
import static scotch.compiler.parser.ScotchLexer.SEMICOLON;
import static scotch.compiler.parser.ScotchLexer.WHERE;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.List;
import org.antlr.v4.runtime.CharStream;
import org.antlr.v4.runtime.Token;
import org.antlr.v4.runtime.TokenFactory;
import org.antlr.v4.runtime.TokenSource;
import org.antlr.v4.runtime.misc.Pair;

public class ScotchLayoutLexer implements TokenSource {

    private static final int LOOK_AHEAD = 5;
    private final ScotchLexer    delegate;
    private final Deque<State>   states;
    private final Deque<Integer> braces;
    private final Deque<Integer> indents;
    private final Deque<Integer> letIndents;
    private final List<Token>    tokens;

    public ScotchLayoutLexer(ScotchLexer delegate) {
        this.delegate = delegate;
        states = new ArrayDeque<>(asList(SCAN_DEFAULT));
        braces = new ArrayDeque<>();
        indents = new ArrayDeque<>();
        letIndents = new ArrayDeque<>();
        tokens = new ArrayList<>();
    }

    @Override
    public int getCharPositionInLine() {
        return firstToken().getCharPositionInLine();
    }

    @Override
    public CharStream getInputStream() {
        return firstToken().getInputStream();
    }

    @Override
    public int getLine() {
        return firstToken().getLine();
    }

    @Override
    public String getSourceName() {
        return firstToken().getTokenSource().getSourceName();
    }

    @Override
    public TokenFactory<?> getTokenFactory() {
        return delegate.getTokenFactory();
    }

    @Override
    public Token nextToken() {
        if (tokens.isEmpty()) {
            buffer();
        } else {
            buffer();
            advance();
        }
        layout();
        return firstToken();
    }

    @Override
    public void setTokenFactory(TokenFactory<?> factory) {
        delegate.setTokenFactory(factory);
    }

    private void accept() {
        enterState(ACCEPT);
    }

    private void advance() {
        tokens.remove(0);
    }

    private void bracesDown() {
        braces.push(braces.pop() - 1);
    }

    private void bracesUp() {
        braces.push(braces.pop() + 1);
    }

    private void buffer() {
        if (tokens.isEmpty() || lastToken().getType() != EOF && tokens.size() < LOOK_AHEAD) {
            buffer_();
        }
    }

    private void buffer_() {
        Token token = delegate.nextToken();
        if (token.getType() == EOF && lastToken().getType() != SEMICOLON) {
            tokens.add(token(SEMICOLON, ";", token));
        }
        tokens.add(token);
        buffer();
    }

    private int currentColumn() {
        return firstToken().getCharPositionInLine();
    }

    private int currentIndent() {
        return indents.peek();
    }

    private int currentLetIndent() {
        return letIndents.peek();
    }

    private State currentState() {
        return states.peek();
    }

    private void dedent() {
        indents.pop();
    }

    private void disableScan() {
        enterState(SCAN_DISABLED);
        braces.push(0);
    }

    private void enterLayout() {
        if (secondToken().getType() == NEWLINE) {
            exciseNewLine();
            enterLayout();
        } else if (secondToken().getType() == OPEN_CURLY) {
            disableScan();
        } else {
            indent(secondToken().getCharPositionInLine());
            insertLCurly();
            enterState(SCAN_LAYOUT);
            accept();
        }
    }

    private void enterLet() {
        if (secondToken().getType() == NEWLINE) {
            exciseNewLine();
            enterLet();
        } else if (secondToken().getType() == OPEN_CURLY) {
            disableScan();
        } else {
            letIndent();
            insertLCurly();
            enterState(SCAN_LET);
            accept();
        }
    }

    private void enterState(State state) {
        states.push(state);
    }

    private void exciseNewLine() {
        if (secondToken().getType() == NEWLINE) {
            tokens.remove(1);
        } else {
            throw new IllegalStateException();
        }
    }

    private Token firstToken() {
        return tokens.get(0);
    }

    private boolean hasBraces() {
        return braces.peek() > 0;
    }

    private boolean in(State state) {
        return currentState() == state;
    }

    private void indent(int level) {
        indents.push(level);
    }

    private void indent() {
        indent(firstToken().getCharPositionInLine());
    }

    private void insertIn() {
        insertToken(token(IN, "in", firstToken()));
    }

    private void insertLCurly() {
        tokens.add(1, token(OPEN_CURLY, "{", tokens.get(2)));
    }

    private void insertRCurly() {
        insertToken(token(CLOSE_CURLY, "}", firstToken()));
    }

    private void insertSemicolon() {
        insertToken(token(SEMICOLON, ";", firstToken()));
    }

    private void insertToken(Token token) {
        tokens.add(0, token);
    }

    private Token lastToken() {
        return tokens.get(tokens.size() - 1);
    }

    private void layout() {
        buffer();
        layout_();
    }

    private void layout_() {
        switch (firstToken().getType()) {
            case EOF:
                if (in(SCAN_LAYOUT)) {
                    leaveLayout();
                }
                return;
            case SEMICOLON:
            case OPEN_SQUARE:
            case CLOSE_SQUARE:
            case OPEN_CURLY:
            case CLOSE_CURLY:
                return;
            case WHERE:
            case ON:
            case DO:
                enterLayout();
                return;
            case LET:
                if (in(SCAN_LET) && currentColumn() == currentLetIndent()) {
                    advance();
                    layout();
                } else {
                    enterLet();
                }
                return;
            case NEWLINE:
                advance();
                layout();
                return;
        }
        switch (currentState()) {
            case ACCEPT:
                leaveState();
                break;
            case SCAN_LAYOUT:
                if (currentColumn() == currentIndent()) {
                    insertSemicolon();
                    accept();
                } else if (currentColumn() < currentIndent()) {
                    leaveLayout();
                }
                break;
            case SCAN_LET:
                if (currentColumn() == currentIndent()) {
                    insertSemicolon();
                    accept();
                } else if (firstToken().getType() == IN) {
                    leaveLet();
                } else if (currentColumn() < currentIndent()) {
                    insertIn();
                    leaveLet();
                }
                break;
            case SCAN_DISABLED:
                if (firstToken().getType() == OPEN_CURLY || firstToken().getType() == OPEN_SQUARE) {
                    bracesUp();
                } else if (firstToken().getType() == CLOSE_CURLY || secondToken().getType() == CLOSE_SQUARE) {
                    if (hasBraces()) {
                        bracesDown();
                    } else {
                        leaveState();
                    }
                }
                break;
            default:
                if (indents.isEmpty()) {
                    indent();
                } else if (currentColumn() <= currentIndent()) {
                    dedent();
                    insertSemicolon();
                    indent(firstToken().getCharPositionInLine());
                    accept();
                }
        }
    }

    private void leaveLayout() {
        leaveState();
        if (in(SCAN_LET)) {
            insertRCurly();
            insertSemicolon();
            dedent();
        } else {
            insertSemicolon();
            insertRCurly();
            insertSemicolon();
            dedent();
            accept();
        }
    }

    private void leaveLet() {
        insertRCurly();
        insertSemicolon();
        letDedent();
        leaveState();
        accept();
    }

    private void leaveState() {
        if (states.isEmpty()) {
            throw new IllegalStateException("Cannot leave root state");
        } else {
            states.pop();
        }
    }

    private void letDedent() {
        letIndents.pop();
        dedent();
    }

    private void letIndent() {
        letIndents.push(firstToken().getCharPositionInLine());
        indent(secondToken().getCharPositionInLine());
    }

    private Token secondToken() {
        return tokens.get(1);
    }

    private Token token(int type, String text, Token token) {
        return getTokenFactory().create(
            new Pair<>(delegate, delegate.getInputStream()), type, text,
            token.getChannel(),
            token.getStartIndex(),
            token.getStopIndex(),
            token.getLine(),
            token.getCharPositionInLine()
        );
    }

    enum State {
        ACCEPT,
        SCAN_DEFAULT,
        SCAN_DISABLED,
        SCAN_LET,
        SCAN_LAYOUT
    }
}