path: root/compiler/codegen.cup

import "compiler/builtins.cup"
import "compiler/ast.cup"
import "std/file.cup"

let gen_out_file: File*;
let gen_label_counter = -1; // So the labels start at 0

let gen_string_literals = vector_new();

fn emit_asm4(msg1: char*, msg2: char*, msg3: char*, msg4: char*) {
    fwrite(gen_out_file, msg1, strlen(msg1));
    fwrite(gen_out_file, msg2, strlen(msg2));
    fwrite(gen_out_file, msg3, strlen(msg3));
    fwrite(gen_out_file, msg4, strlen(msg4));
}

fn emit_asm3(msg1: char*, msg2: char*, msg3: char*) {
    fwrite(gen_out_file, msg1, strlen(msg1));
    fwrite(gen_out_file, msg2, strlen(msg2));
    fwrite(gen_out_file, msg3, strlen(msg3));
}

fn emit_asm2(msg1: char*, msg2: char*) {
    fwrite(gen_out_file, msg1, strlen(msg1));
    fwrite(gen_out_file, msg2, strlen(msg2));
}

fn emit_asm(msg: char*) {
    fwrite(gen_out_file, msg, strlen(msg));
}

// This is a hack because we don't support manual escaping, and 
// we need to be able to include '`' in string literal definitions.
fn emit_asm_char(c: char) {
    fwrite(gen_out_file, &c, 1);
}

fn emit_num(num: int) {
    // FIXME: Just support printing negatives directly.
    if (num < 0) {
        emit_asm("-");
        num = -num;
    }
    fputu(gen_out_file, num);
}

fn generate_syscall(num: int) {
    emit_asm("    mov rax, "); emit_num(num); emit_asm("\n");
    emit_asm("    syscall\n");
}

fn subregister_for_type(typ: Type*): char* {
    let n = size_for_type(typ);
    if (n == 1) return "al";
    if (n == 2) return "ax";
    if (n == 4) return "eax";
    if (n == 8) return "rax";
    die2(here, "Unsupported type size");
}

fn specifier_for_type(typ: Type*): char* {
    let n = size_for_type(typ);
    if (n == 1) return "byte";
    if (n == 2) return "word";
    if (n == 4) return "dword";
    if (n == 8) return "qword";
    die2(here, "Unsupported type size");
}

fn generate_expr_into_rax(node: Node*);

fn generate_lvalue_into_rax(node: Node*) {
    if (node.typ == AST_LOCAL_VAR) {
        let offset = node.d.variable.offset;
        emit_asm("    mov rax, rbp\n");
        emit_asm("    sub rax, "); emit_num(offset); emit_asm("\n");
    } else if (node.typ == AST_GLOBAL_VAR) {
        let offset = node.d.variable.offset;
        emit_asm("    mov rax, global_vars\n");
        emit_asm("    add rax, "); emit_num(offset); emit_asm("\n");
    } else if (node.typ == AST_MEMBER) {
        let offset = node.d.member.offset;
        if (node.d.member.is_ptr)
            generate_expr_into_rax(node.d.member.obj);
        else
            generate_lvalue_into_rax(node.d.member.obj);
        emit_asm("    add rax, "); emit_num(offset); emit_asm("\n");
    } else if (node.typ == AST_DEREF) {
        generate_expr_into_rax(node.d.unary);
    } else {
        die2("Unsupported type in generate_lvalue_into_rax: ", node_type_to_string(node.typ));
    }
}

fn generate_function_call(node: Node*) {
    let total_size = 0;
    let n = node.d.call.args.size;
    for (let i = n-1; i >= 0; --i) {
        let expr: Node* = node.d.call.args.data[i];
        generate_expr_into_rax(expr);
        emit_asm("    push rax\n");
        // TODO: this might be an issue if we pass structs some day
        total_size = total_size + 8;
    }
    emit_asm3("    call func_", node.d.call.func.d.func.name, "\n");
    emit_asm("    add rsp, "); emit_num(total_size); emit_asm("\n");
}

fn generate_expr_into_rax(node: Node*) {
    if (node.typ == AST_LITERAL) {
        if (node.etyp.typ == TYPE_INT) {
            emit_asm("    mov rax, "); emit_num(node.d.literal.as_int); emit_asm("\n");
        } else if (node.etyp.typ == TYPE_CHAR) {
            emit_asm("    mov rax, "); emit_num(node.d.literal.as_char); emit_asm("\n");
        } else if (node.etyp.typ == TYPE_PTR) {
            let s = node.d.literal.as_string;
            let idx = gen_string_literals.size;
            vector_push(gen_string_literals, s);
            emit_asm("    mov rax, global_string_"); emit_num(idx); emit_asm("\n");
        } else {
            die("Unsupported literal type in generate_expr_into_rax");
        }
    } else if (node.typ == AST_ADDROF) {
        generate_lvalue_into_rax(node.d.unary);

    } else if (node.typ == AST_CONDITIONAL) {
        let label = ++gen_label_counter;
        generate_expr_into_rax(node.d.conditional.cond);
        emit_asm("    cmp rax, 0\n");
        emit_asm("    je .cond_else_"); emit_num(label); emit_asm("\n");
        generate_expr_into_rax(node.d.conditional.then);
        emit_asm("    jmp .cond_end_"); emit_num(label); emit_asm("\n");
        emit_asm(".cond_else_"); emit_num(label); emit_asm(":\n");
        generate_expr_into_rax(node.d.conditional.els);
        emit_asm(".cond_end_"); emit_num(label); emit_asm(":\n");

    } else if (node.typ == AST_PLUS) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    add rax, rbx\n");

    } else if (node.typ == AST_MINUS) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    sub rax, rbx\n");

    } else if (node.typ == AST_DIV) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cqo\n");
        emit_asm("    idiv rbx\n");

    } else if (node.typ == AST_MOD) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cqo\n");
        emit_asm("    idiv rbx\n");
        emit_asm("    mov rax, rdx\n");

    } else if (node.typ == AST_MUL) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    imul rbx\n");

    } else if (node.typ == AST_EQ) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cmp rax, rbx\n");
        emit_asm("    sete al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_NEQ) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cmp rax, rbx\n");
        emit_asm("    setne al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_LT) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cmp rax, rbx\n");
        emit_asm("    setl al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_LEQ) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cmp rax, rbx\n");
        emit_asm("    setle al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_GT) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cmp rax, rbx\n");
        emit_asm("    setg al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_GEQ) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cmp rax, rbx\n");
        emit_asm("    setge al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_BWAND) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    and rax, rbx\n");

    } else if (node.typ == AST_BWOR) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    or rax, rbx\n");

    } else if (node.typ == AST_XOR) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    xor rax, rbx\n");

    } else if (node.typ == AST_BWINV) {
        generate_expr_into_rax(node.d.unary);
        emit_asm("    not rax\n");

    } else if (node.typ == AST_NEG) {
        generate_expr_into_rax(node.d.unary);
        emit_asm("    neg rax\n");

    } else if (node.typ == AST_NOT) {
        generate_expr_into_rax(node.d.unary);
        emit_asm("    cmp rax, 0\n");
        emit_asm("    sete al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_OR) {
        // With short circuiting!
        let label = ++gen_label_counter;
        generate_expr_into_rax(node.d.binary.lhs,);
        // If left is true, we can short-circuit
        emit_asm("    cmp rax, 0\n");
        emit_asm("    je .or_right_"); emit_num(label); emit_asm("\n");
        emit_asm("    mov rax, 1\n");
        emit_asm("    jmp .or_end_"); emit_num(label); emit_asm("\n");
        emit_asm(".or_right_"); emit_num(label); emit_asm(":\n");
        generate_expr_into_rax(node.d.binary.rhs);
        // Booleanize the result
        emit_asm("    cmp rax, 0\n");
        emit_asm("    setne al\n");
        emit_asm(".or_end_"); emit_num(label); emit_asm(":\n");

    } else if (node.typ == AST_AND) {
        let label = ++gen_label_counter;
        generate_expr_into_rax(node.d.binary.lhs);
        // If left is false, we can short-circuit
        emit_asm("    cmp rax, 0\n");
        emit_asm("    jne .and_right_"); emit_num(label); emit_asm("\n");
        emit_asm("    mov rax, 0\n");
        emit_asm("    jmp .and_end_"); emit_num(label); emit_asm("\n");
        emit_asm(".and_right_"); emit_num(label); emit_asm(":\n");
        generate_expr_into_rax(node.d.binary.rhs);
        // Booleanize the result
        emit_asm("    cmp rax, 0\n");
        emit_asm("    setne al\n");
        emit_asm(".and_end_"); emit_num(label); emit_asm(":\n");

    } else if (is_lvalue(node.typ)) {
        generate_lvalue_into_rax(node);
        if (size_for_type(node.etyp) == 8) {
            emit_asm("    mov rax, [rax]\n");
        } else {
            emit_asm3("    movsx rax, ", specifier_for_type(node.etyp), " [rax]\n");
        }

    } else if (node.typ == AST_ASSIGN) {
        let var = node.d.assign.lhs;
        generate_lvalue_into_rax(var);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.assign.rhs);
        emit_asm("    pop rbx\n");
        emit_asm3("    mov [rbx], ", subregister_for_type(var.etyp), "\n");

    } else if (node.typ == AST_FUNCCALL) {
        generate_function_call(node);

    } else {
        die2("Unsupported node type in generate_expr_into_rax: ", node_type_to_string(node.typ));
    }
}

fn generate_block(node: Node*);

fn generate_statement(node: Node*) {
    if (node.typ == AST_RETURN) {
        if (node.d.unary) {
            generate_expr_into_rax(node.d.unary);
        } else {
            emit_asm("    xor rax, rax\n"); // Default to 0
        }

        emit_asm("    push rax\n");
        // TODO: Undo the defer stack here, this is for consistency with the C implementation for now.
        emit_asm("    pop rax\n");

        emit_asm("    mov rsp, rbp\n");
        emit_asm("    pop rbp\n");
        emit_asm("    ret\n");
    } else if (node.typ == AST_VARDECL) {
        if (node.d.var_decl.init) {
            generate_expr_into_rax(node.d.var_decl.init);
            let offset = node.d.var_decl.var.offset;
            emit_asm("    mov [rbp-"); emit_num(offset); emit_asm("], rax\n");
        }
    } else if (node.typ == AST_BLOCK) {
        generate_block(node);
    } else if (node.typ == AST_IF) {
        let label = ++gen_label_counter;
        generate_expr_into_rax(node.d.conditional.cond);
        // If we don't have an `else` clause, we can simplify
        if (node.d.conditional.els == null) {
            emit_asm("    cmp rax, 0\n");
            emit_asm("    je .if_end_"); emit_num(label); emit_asm("\n");
            generate_statement(node.d.conditional.then);
            emit_asm(".if_end_"); emit_num(label); emit_asm(":\n");
        } else {
            emit_asm("    cmp rax, 0\n");
            emit_asm("    je .if_else_"); emit_num(label); emit_asm("\n");
            generate_statement(node.d.conditional.then);
            emit_asm("    jmp .if_end_"); emit_num(label); emit_asm("\n");
            emit_asm(".if_else_"); emit_num(label); emit_asm(":\n");
            generate_statement(node.d.conditional.els);
            emit_asm(".if_end_"); emit_num(label); emit_asm(":\n");
        }
    } else if (node.typ == AST_WHILE) {
        let label = ++gen_label_counter;
        emit_asm(".loop_start_"); emit_num(label); emit_asm(":\n");
        emit_asm(".loop_continue_"); emit_num(label); emit_asm(":\n");
        generate_expr_into_rax(node.d.looop.cond);
        emit_asm("    cmp rax, 0\n");
        emit_asm("    je .loop_end_"); emit_num(label); emit_asm("\n");
        generate_statement(node.d.looop.body);
        emit_asm("    jmp .loop_start_"); emit_num(label); emit_asm("\n");
        emit_asm(".loop_end_"); emit_num(label); emit_asm(":\n");

    } else if (node.typ == AST_FOR) {
        let label = ++gen_label_counter;
        if (node.d.looop.init)
            generate_statement(node.d.looop.init);
        emit_asm(".loop_start_"); emit_num(label); emit_asm(":\n");
        if (node.d.looop.cond) {
            generate_expr_into_rax(node.d.looop.cond);
            emit_asm("    cmp rax, 0\n");
            emit_asm("    je .loop_end_"); emit_num(label); emit_asm("\n");
        }
        generate_statement(node.d.looop.body);
        emit_asm(".loop_continue_"); emit_num(label); emit_asm(":\n");
        if (node.d.looop.step)
            generate_statement(node.d.looop.step);
        emit_asm("    jmp .loop_start_"); emit_num(label); emit_asm("\n");
        emit_asm(".loop_end_"); emit_num(label); emit_asm(":\n");

    } else {
        // Default to a simple expression statement
        generate_expr_into_rax(node);
    }
}

fn generate_block(node: Node*) {
    let n = node.d.block.children.size;
    for (let i = 0; i < n; ++i) {
        generate_statement(node.d.block.children.data[i]);
    }
}

fn generate_function(node: Node*) {
    // Skip declarations
    if (node.d.func.body == null)
        return;

    emit_asm3("global func_", node.d.func.name, "\n");
    emit_asm3("func_", node.d.func.name, ":\n");
    emit_asm("    push rbp\n");
    emit_asm("    mov rbp, rsp\n");
    emit_asm("    sub rsp, "); emit_num(node.d.func.max_locals_size); emit_asm("\n");

    generate_block(node.d.func.body);

    emit_asm("    mov rsp, rbp\n");
    emit_asm("    pop rbp\n");
    // Return 0 by default if we don't have a return statement
    emit_asm("    mov qword rax, 0\n");
    emit_asm("    ret\n");
}

fn generate_program(ast: Node*, file: File*) {
    gen_out_file = file;

    let n = ast.d.block.children.size;
    for (let i = 0; i < n; ++i) {
        let node: Node* = ast.d.block.children.data[i];
        if (node.typ == AST_FUNC) {
            generate_function(node);
        } else if (node.typ == AST_VARDECL) {
            // Do nothing
        } else {
            die("Unknown node type in generate_program");
        }
    }

    if (OS_IS_MACOS) {
        emit_asm("global _main\n");
        emit_asm("_main:\n");
        // Push argv
        emit_asm("    mov rax, rsi\n");
        emit_asm("    push rax\n");
        // Push argc
        emit_asm("    mov rax, rdi\n");
        emit_asm("    push rax\n");
    } else {
        emit_asm("global _start\n");
        emit_asm("_start:\n");

        emit_asm("    mov rbp, rsp\n");
        // // Push argv
        emit_asm("    mov rax, rbp\n");
        emit_asm("    add rax, 8\n");
        emit_asm("    push rax\n");
        // Push argc
        emit_asm("    mov rax, [rbp]\n");
        emit_asm("    push rax\n");
    }

    // Initialize all the global variables
    for (let i = 0; i < n; ++i) {
        let node: Node* = ast.d.block.children.data[i];
        if (node.typ == AST_VARDECL && node.d.var_decl.init) {
            let expr = node.d.var_decl.init;
            generate_expr_into_rax(node.d.var_decl.init);
            let offset = node.d.var_decl.var.offset;
            emit_asm("    mov rbx, global_vars\n");
            emit_asm("    add rbx, "); emit_num(offset); emit_asm("\n");
            emit_asm3("    mov [rbx], ", subregister_for_type(expr.etyp), "\n");
        }
    }

    emit_asm("    call func_main\n");
    emit_asm("    mov rdi, rax\n");
    generate_syscall(SYS_exit);

    generate_builtins();

    emit_asm("section .bss\n");
    emit_asm("    global_vars: resb "); emit_num(p_global_offset); emit_asm("\n");

    // Global strings
    emit_asm("section .data\n");
    // TODO: Don't to this here because a string containing a backtick will
    //       cause invalid output and break everything. Maybe just output the
    //       byte values.
    for (let i = 0; i < gen_string_literals.size; ++i) {
        emit_asm("    global_string_"); emit_num(i); emit_asm(": db ");
        emit_asm_char('`');
        emit_asm(gen_string_literals.data[i]);
        emit_asm_char('`');
        emit_asm(", 0\n");
    }
}