gombaila/src/main/java/ch/fritteli/gombaila/domain/generator/Generator.java

package ch.fritteli.gombaila.domain.generator;

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 org.jetbrains.annotations.NotNull;

import java.util.HashMap;
import java.util.Map;

public class Generator {
    @NotNull
    private static final String header = """
            global _start
            _start:
            """;
    @NotNull
    private static final String defaultExitStmt = """
                mov rax, 60
                mov rdi, 0
                syscall
            """;

    @NotNull
    private final NodeProg nodeProg;
    @NotNull
    private final StringBuilder asm;
    @NotNull
    private final ExprVisitor exprVisitor;
    @NotNull
    private final StmtVisitor stmtVisitor;
    @NotNull
    private final Map<String, Variable> identifierStackposition = new HashMap<>();
    private int stackSize = 0;

    public Generator(@NotNull final NodeProg nodeProg) {
        this.nodeProg = nodeProg;
        this.asm = new StringBuilder(header);
        this.exprVisitor = new ExprVisitor();
        this.stmtVisitor = new StmtVisitor();
    }

    @NotNull
    public String generate() {
        for (final NodeStmt stmt : this.nodeProg.stmts()) {
            this.generateStmt(stmt);
        }
        this.asm.append(defaultExitStmt);
        return this.asm.toString();
    }

    private void generateStmt(@NotNull final NodeStmt stmt) {
        this.stmtVisitor.visit(stmt);
    }

    private void generateExpr(@NotNull final NodeExpr expr) {
        this.exprVisitor.visit(expr);
    }

    private void push(@NotNull final String reg) {
        this.line("push %s", reg);
        this.stackSize++;
    }

    private void pop(@NotNull final String reg) {
        this.line("pop %s", reg);
        this.stackSize--;
    }

    private void label(@NotNull final String name) {
        this.asm.append(name)
                .append(":")
                .append("\n");
    }

    private void line(@NotNull final String line, @NotNull final Object... params) {
        this.asm.append("    ")
                .append(line.formatted(params))
                .append("\n");
    }

    private record Variable(int stackLocation) {
    }

    private final class StmtVisitor {

        private StmtVisitor() {
        }

        void visit(@NotNull final NodeStmt stmt) {
            switch (stmt) {
                case final NodeStmtLet stmtLet -> visit(stmtLet);
                case final NodeStmtExit stmtExit -> visit(stmtExit);
                case final NodeStmtPrint stmtPrint -> visit(stmtPrint);
            }
        }

        private void visit(@NotNull final NodeStmtExit stmt) {
            generateExpr(stmt.expr());
            line("mov rax, 60");
            pop("rdi");
            line("syscall");
        }

        private void visit(@NotNull final NodeStmtLet stmt) {
            if (identifierStackposition.containsKey(stmt.ident().value().get())) {
                throw new GeneratorException("Identifier already used", stmt);
            }
            identifierStackposition.put(((String) stmt.ident().value().get()), new Variable(stackSize));
            generateExpr(stmt.expr());
        }

        private void visit(@NotNull final NodeStmtPrint stmt) {
            generateExpr(stmt.nodeExpr());
            // rdx: length of output
            line("mov rdx, 8");
            // rsi: output
            pop("rsi");
            // rdi: fd-value (STDOUT=1)
            line("mov rdi, 1");
            // rax: syscall (SYS_WRITE=1)
            line("mov rax, 1");
            // syscall
            line("syscall");
        }
    }

    private final class ExprVisitor {
        private ExprVisitor() {
        }

        void visit(@NotNull final NodeExpr expr) {
            switch (expr) {
                case final NodeExprIdent exprIdent -> visit(exprIdent);
                case final NodeExprIntLit exprIntLit -> visit(exprIntLit);
                case final NodeBinExprAdd exprAdd -> visit(exprAdd);
                case final NodeBinExprMinus exprMinus -> visit(exprMinus);
                case final NodeBinExprMult exprMult -> visit(exprMult);
                case final NodeBinExprDiv exprDiv -> visit(exprDiv);
                case final NodeBinExprMod exprMod -> visit(exprMod);
                case final NodeBinExprExp exprExp -> visit(exprExp);
            }
        }

        private void visit(@NotNull final NodeExprIdent expr) {
            final Variable variable = identifierStackposition.get(expr.ident().value().get());
            if (variable == null) {
                throw new GeneratorException("Undeclared identifier", expr);
            }
            push("QWORD [rsp+%d]".formatted(8 * (stackSize - variable.stackLocation - 1)));
        }

        private void visit(@NotNull final NodeExprIntLit expr) {
            line("mov rax, %s", expr.value().value().get());
            push("rax");
        }

        private void visit(@NotNull final NodeBinExprAdd expr) {
            generateExpr(expr.lhs());
            generateExpr(expr.rhs());
            pop("rbx");
            pop("rax");
            line("add rax, rbx");
            push("rax");
        }

        private void visit(@NotNull final NodeBinExprMinus expr) {
            generateExpr(expr.lhs());
            generateExpr(expr.rhs());
            pop("rbx");
            pop("rax");
            line("sub rax, rbx");
            push("rax");
        }

        private void visit(@NotNull final NodeBinExprMult expr) {
            throw new GeneratorException("Not yet implemented!", expr);
        }

        private void visit(@NotNull final NodeBinExprDiv expr) {
            throw new GeneratorException("Not yet implemented!", expr);
        }

        private void visit(@NotNull final NodeBinExprMod expr) {
            throw new GeneratorException("Not yet implemented!", expr);
        }

        private void visit(@NotNull final NodeBinExprExp expr) {
            throw new GeneratorException("Not yet implemented!", expr);
        }
    }
}