gombaila/src/main/java/ch/fritteli/gombaila/domain/parser/Parser.java

package ch.fritteli.gombaila.domain.parser;

import ch.fritteli.gombaila.domain.SeqWalker;
import ch.fritteli.gombaila.domain.common.NodeBinExprAdd;
import ch.fritteli.gombaila.domain.common.NodeBinExprDiv;
import ch.fritteli.gombaila.domain.common.NodeBinExprExp;
import ch.fritteli.gombaila.domain.common.NodeBinExprMinus;
import ch.fritteli.gombaila.domain.common.NodeBinExprMod;
import ch.fritteli.gombaila.domain.common.NodeBinExprMult;
import ch.fritteli.gombaila.domain.common.NodeExpr;
import ch.fritteli.gombaila.domain.common.NodeExprIdent;
import ch.fritteli.gombaila.domain.common.NodeExprIntLit;
import ch.fritteli.gombaila.domain.common.NodeProg;
import ch.fritteli.gombaila.domain.common.NodeStmt;
import ch.fritteli.gombaila.domain.common.NodeStmtExit;
import ch.fritteli.gombaila.domain.common.NodeStmtLet;
import ch.fritteli.gombaila.domain.common.NodeStmtPrint;
import ch.fritteli.gombaila.domain.common.Token;
import ch.fritteli.gombaila.domain.common.TokenType;
import io.vavr.collection.List;
import io.vavr.collection.Seq;
import org.jetbrains.annotations.NotNull;

public class Parser {
    @NotNull
    private final SeqWalker<Token> tokens;

    public Parser(@NotNull final Seq<Token> tokens) {

        this.tokens = new SeqWalker<>(tokens.reject(token -> TokenType.WHITESPACE.equals(token.type())));
    }

    @NotNull
    public NodeProg parse() {
        Seq<NodeStmt> stmts = List.empty();
        while (this.tokens.hasNext()) {
            stmts = stmts.append(this.parseStmt());
        }

        return new NodeProg(stmts);
    }

    @NotNull
    private NodeStmt parseStmt() {
        final NodeStmt result;
        if (this.checkNextTokenTypeConsuming(TokenType.EXIT)) {
            this.assertAndConsumeNextTokenType(TokenType.OPEN_PAREN);
            final NodeExpr nodeExpr = this.parseExpr(1);
            this.assertAndConsumeNextTokenType(TokenType.CLOSE_PAREN);
            result = new NodeStmtExit(nodeExpr);
        } else if (this.checkNextTokenTypeConsuming(TokenType.LET)) {
            final Token identifier = this.assertAndConsumeNextTokenType(TokenType.IDENTIFIER);
            this.assertAndConsumeNextTokenType(TokenType.EQUALS);
            final NodeExpr nodeExpr = this.parseExpr(1);
            result = new NodeStmtLet(identifier, nodeExpr);
        } else if (this.checkNextTokenTypeConsuming(TokenType.PRINT)) {
            this.assertAndConsumeNextTokenType(TokenType.OPEN_PAREN);
            final NodeExpr nodeExpr = this.parseExpr(1);
            this.assertAndConsumeNextTokenType(TokenType.CLOSE_PAREN);
            result = new NodeStmtPrint(nodeExpr);
        } else {
            throw new ParserException("Could not parse statement", null);
        }

        this.assertAndConsumeNextTokenType(TokenType.SEMI);
        return result;
    }

    @NotNull
    private NodeExpr parseExpr(final int minPrecedence) {
        NodeExpr result;
        if (this.checkNextTokenType(TokenType.INT_LIT)) {
            result = new NodeExprIntLit(this.tokens.next());
        } else if (this.checkNextTokenType(TokenType.IDENTIFIER)) {
            result = new NodeExprIdent(this.tokens.next());
        } else if (this.checkNextTokenTypeConsuming(TokenType.OPEN_PAREN)) {
            result = this.parseExpr(1);
            this.assertAndConsumeNextTokenType(TokenType.CLOSE_PAREN);
        } else {
            throw new ParserException(null, this.tokens.peekNext().getOrNull());
        }

        while (this.hasNextTokenPrecedenceGTE(minPrecedence)) {
            final Token token = this.tokens.next();
            final int precedence = token.type().precedence().get();
            final TokenType.Associativity associativity = token.type().associativity().get();
            final int nextMinPrecedence;
            if (associativity == TokenType.Associativity.LEFT) {
                nextMinPrecedence = precedence + 1;
            } else {
                nextMinPrecedence = precedence;
            }
            final NodeExpr rhs = this.parseExpr(nextMinPrecedence);
            result = switch (token.type()) {
                case PLUS -> new NodeBinExprAdd(result, rhs);
                case MINUS -> new NodeBinExprMinus(result, rhs);
                case MULT -> new NodeBinExprMult(result, rhs);
                case DIV -> new NodeBinExprDiv(result, rhs);
                case MOD -> new NodeBinExprMod(result, rhs);
                case EXP -> new NodeBinExprExp(result, rhs);
                default -> throw new ParserException("Expected binary operator token", token);
            };
        }

        return result;
    }

    private boolean hasNextTokenPrecedenceGTE(final int minPrecedence) {
        return this.tokens.peekNext()
                .map(Token::type)
                .flatMap(TokenType::precedence)
                .exists(precedence -> precedence >= minPrecedence);
    }

    private boolean checkNextTokenType(@NotNull final TokenType type) {
        return this.tokens.peekNext()
                .exists(token -> token.type().equals(type));
    }

    private boolean checkNextTokenTypeConsuming(@NotNull final TokenType type) {
        if (this.tokens.peekNext()
                .exists(token -> token.type().equals(type))) {
            this.tokens.next();
            return true;
        }
        return false;
    }

    @NotNull
    private Token assertAndConsumeNextTokenType(@NotNull final TokenType type) {
        if (this.checkNextTokenType(type)) {
            return this.tokens.next();
        }
        throw new ParserException("Unexpected token", this.tokens.peekNext().getOrNull());
    }
}