codegen.cpp

/**
 * codegen.cpp
 * author: Bao Le
 */

#include "codegen.hpp"
#include <stdexcept>
#include <vector>

// Entry point
void CodeGenerator::generate(const std::string& filename) {
    output.open(filename);
    if (!output.is_open()) {
        throw std::runtime_error("[CODEGEN] Couldn't open filename\n");
    }
    // text section for assembly
    emit("  .text");
    for (auto& child : root->children) {
        if (auto* fn = dynamic_cast<Function*>(child.get())) {
            emit_function(fn);
        }
    }
    emit_string_literals_section();
    output.close();
}

std::string CodeGenerator::escape_asm_string(const std::string& s) {
    std::string out;
    for (char c : s) {
        if (c == '\\') out += "\\\\";
        else if (c == '"') out += "\\\"";
        else if (c == '\n') out += "\\n";
        else if (c == '\t') out += "\\t";
        else out += c;
    }
    return out;
}

void CodeGenerator::emit_string_literals_section() {
    if (string_literals.empty()) return;
    emit("");
    emit(".section __TEXT,__cstring,cstring_literals");
    for (auto& entry : string_literals) {
        emit(entry.first + ":");
        emit("    .asciz \"" + escape_asm_string(entry.second) + "\"");
    }
}

// Helper function to write to assembly file
void CodeGenerator::emit(const std::string& line) {
    output << line << '\n';
}

// Specific emitters (write to assembly)
void CodeGenerator::emit_function(Function* fn) {
    stack_offsets.clear();
    current_offset = 0;

    // scan to calculate stack size
    int locals = count_locals(fn->body.get());
    int params = fn->params.size();
    int spill_bytes = 64; // scratch space for up to 8 call arguments
    int total_bytes = (locals + params) * 8 + 16 + spill_bytes;
    // round up stack size
    int stack_size = (total_bytes + 15) & ~15; // flip the last 4 bits to 0
    current_stack_size = stack_size;

    // prologue
    emit(".section __TEXT,__text");
    emit(".globl _" + fn->name);
    emit("_" + fn->name + ":");
    emit("    sub sp, sp, #" + std::to_string(stack_size)); // allocate stack size
    emit("    stp x29, x30, [sp, #" + std::to_string(stack_size - 16) + "]"); // store the current frame pointer and link register from sp + (stack_size - 16)
    emit("    add x29, sp, #" + std::to_string(stack_size - 16)); // move x29 to the point to same address that holds the caller's x29 (frame pointer) for function call
    int epilogue_label = label_count++;
    current_epilogue_label = epilogue_label;

    // store params from registers onto stack
    for (int i = 0; i < fn->params.size(); i++) {
        current_offset -= 8;
        // store this like {'a': -8} so the stack can reference back later
        stack_offsets[fn->params[i].name] = current_offset;
        // store from xN to memory at x29 + current_offset
        emit("    str x" + std::to_string(i) + ", [x29, #" + std::to_string(current_offset) + "]");
    }

    // emit body
    for (auto& stmt : fn->body->statements) {
        emit_statement(stmt.get(), fn->type);
    }

    // epilogue (fallthrough path)
    emit(".Lepilogue_" + std::to_string(epilogue_label) + ":");
    emit_epilogue();
}

void CodeGenerator::emit_epilogue() {
    emit("    ldp x29, x30, [sp, #" + std::to_string(current_stack_size - 16) + "]");
    emit("    add sp, sp, #" + std::to_string(current_stack_size));
    emit("    ret");
}

void CodeGenerator::emit_statement(ASTNode* stmt, ConstantType return_type) {
    if (auto* decl = dynamic_cast<Declaration*>(stmt)) {
        emit_declaration(decl);
        return;
    }
    if (auto* assign = dynamic_cast<Assignment*>(stmt)) {
        emit_assignment(assign);
        return;
    }
    if (auto* if_stmt = dynamic_cast<If*>(stmt)) {
        emit_if(if_stmt, return_type);
        return;
    }
    if (auto* ret = dynamic_cast<Return*>(stmt)) {
        emit_return(ret);
        return;
    }
    if (auto* fn_call = dynamic_cast<FunctionCall*>(stmt)) {
        emit_function_call(fn_call);
        return;
    }
    if (auto* body = dynamic_cast<StatementBody*>(stmt)) {
        for (auto& nested : body->statements) {
            emit_statement(nested.get(), return_type);
        }
        return;
    }
    if (auto* while_stmt = dynamic_cast<While*>(stmt)) {
        emit_while(while_stmt);
        return;
    }
    throw std::runtime_error("[STATEMENT] Unable to dispatch based on node type\n");
}

void CodeGenerator::emit_declaration(Declaration* decl) {
    // assign stack offset
    current_offset -= 8;
    stack_offsets[decl->var.name] = current_offset;
    // if there is an initialization value
    if (decl->init_val) {
        emit_expression(decl->init_val->get());
        // store the final result in x0 
        emit("    str x0, [x29, #" + std::to_string(current_offset) + "]");
    }   
    // if there is no initialization value
    else {
        emit("    str xzr, [x29, #" + std::to_string(current_offset) + "]"); // store 0 in the memory location
    }

}
void CodeGenerator::emit_assignment(Assignment* assign) {
    auto* target_var = dynamic_cast<Variable*>(assign->target.get());
    if (!target_var) {
        throw std::runtime_error("[ASSIGNMENT] target must be a variable\n");
    }
    int offset = stack_offsets.at(target_var->name);
    emit_expression(assign->rhs.get());
    emit("    str x0, [x29, #" + std::to_string(offset) + "]");
}

void CodeGenerator::emit_if(If* if_stmt, ConstantType return_type) {
    // result lands in x0 (0=false, 1=true)
    int else_label = label_count++;
    int end_label = label_count++;

    emit_expression(if_stmt->cond.get());
    emit("    cmp x0, #0");
    // Branch to else if condition is 0 (false)
    emit("    beq .Lelse_" + std::to_string(else_label));
    // emit if_body statements if the condition is true
    for (auto& stmt : if_stmt->if_body->statements) {
        emit_statement(stmt.get(), return_type);
    }
    // jump over else because the if block was executed (true)
    emit("    b .Lend_" + std::to_string(end_label));
    // else label 
    emit(".Lelse_" + std::to_string(else_label) + ":");
    // emit else_body statements (if present)
    if (if_stmt->else_body.has_value()) {
        for (auto& stmt : if_stmt->else_body.value()->statements) {
            emit_statement(stmt.get(), return_type);
        }
    }
    // end label 
    emit(".Lend_" + std::to_string(end_label) + ":");
}

void CodeGenerator::emit_while(While* while_stmt) {
    int head = label_count++;
    int end = label_count++;
    emit(".Lwhile_head_" + std::to_string(head) + ":");
    emit_expression(while_stmt->cond.get());
    emit("    cmp x0, #0");
    emit("    beq .Lwhile_end_" + std::to_string(end));
    for (auto& stmt : while_stmt->body->statements) {
        emit_statement(stmt.get(), ConstantType::Void);
    }
    emit("    b .Lwhile_head_" + std::to_string(head));
    emit(".Lwhile_end_" + std::to_string(end) + ":");
}

void CodeGenerator::emit_return(Return* ret) {
    if (ret->val.has_value()) {
        emit_expression(ret->val.value().get());
    }
    emit("    b .Lepilogue_" + std::to_string(current_epilogue_label));
}

void CodeGenerator::emit_expression(ASTNode* expr) {
    if (auto* constant = dynamic_cast<Constant*>(expr)) {
        if (constant->type == ConstantType::String) {
            std::string label = ".Lstr_" + std::to_string(string_label_count++);
            string_literals.emplace_back(label, constant->val);
            // find the address of the string literal
            emit("    adrp x0, " + label + "@PAGE");
            // add the page offset to the address to reach the first letter of the string
            emit("    add x0, x0, " + label + "@PAGEOFF");
            return;
        }
        std::string imm = constant->val;
        if (constant->type == ConstantType::Bool) {
            imm = (constant->val == "true") ? "1" : "0";
        } else if (constant->type == ConstantType::Char) {
            // store the char as an unsigned integer
            imm = std::to_string(static_cast<unsigned char>(constant->val[0]));
        }
        emit("    mov x0, #" + imm);
        return;
    }
    if (auto* var = dynamic_cast<Variable*>(expr)) {
        int offset = stack_offsets.at(var->name);
        emit("    ldr x0, [x29, #" + std::to_string(offset) + "]");
        return;
    }
    if (auto* bin_op = dynamic_cast<BinaryOperation*>(expr)) {
        emit_expression(bin_op->left.get());
        emit("    mov x1, x0");
        emit_expression(bin_op->right.get());
        // x0 = right, x1 = left
        switch (bin_op->op) {
            case BinaryOperator::Plus:
                emit("    add x0, x1, x0");
                break;
            case BinaryOperator::Minus:
                emit("    sub x0, x1, x0");
                break;
            case BinaryOperator::Star:
                emit("    mul x0, x1, x0");
                break;
            case BinaryOperator::Slash:
                emit("    sdiv x0, x1, x0");
                break;
            case BinaryOperator::LessThan:
                emit("    cmp x1, x0");
                // store the result in x0
                emit("    cset x0, lt");
                break;
            case BinaryOperator::GreaterThan:
                emit("    cmp x1, x0");
                emit("    cset x0, gt");
                break;
            case BinaryOperator::Equal:
                emit("    cmp x1, x0");
                emit("    cset x0, eq");
                break;
            case BinaryOperator::NotEqual:
                emit("    cmp x1, x0");
                emit("    cset x0, ne");
                break;
            default:
                throw std::runtime_error("[EXPRESSION] Unknown binary operator\n");
        }
        return;
    }
    if (auto* fn_call = dynamic_cast<FunctionCall*>(expr)) {
        emit_function_call(fn_call);
        return;
    }
    throw std::runtime_error("[EXPRESSION] Unable to dispatch based on node type\n");
}

void CodeGenerator::emit_function_call(FunctionCall* call) {
    int n = static_cast<int>(call->args.size());

    // Apple ARM64 variadic ABI (printf): x0 = format string, extra args on stack at sp
    if (is_variadic_libc_call(call->name)) {
        emit_expression(call->args[0].get());
        int fmt_slot = current_offset - 8;
        emit("    str x0, [x29, #" + std::to_string(fmt_slot) + "]");

        int varargs = n - 1;
        int stack_bytes = varargs > 0 ? ((varargs * 8 + 15) & ~15) : 16;
        emit("    sub sp, sp, #" + std::to_string(stack_bytes));

        for (int i = 1; i < n; i++) {
            emit_expression(call->args[i].get());
            emit("    str x0, [sp, #" + std::to_string(8 * (i - 1)) + "]");
        }

        emit("    ldr x0, [x29, #" + std::to_string(fmt_slot) + "]");
        emit("    mov x8, #0");
        emit("    bl _" + call->name);
        emit("    add sp, sp, #" + std::to_string(stack_bytes));
        return;
    }

    std::vector<int> spill_offsets;
    int spill_offset = current_offset;

    for (int i = 0; i < n; i++) {
        emit_expression(call->args[i].get());
        spill_offset -= 8;
        spill_offsets.push_back(spill_offset);
        emit("    str x0, [x29, #" + std::to_string(spill_offset) + "]");
    }
    for (int i = 0; i < n; i++) {
        emit("    ldr x" + std::to_string(i)
                + ", [x29, #" + std::to_string(spill_offsets[i]) + "]");
    }
    emit("    bl _" + call->name);
}

bool CodeGenerator::is_variadic_libc_call(const std::string& name) const {
    return name == "printf";
}

int CodeGenerator::count_locals(StatementBody* body) {
    int decl_count {};
    // Declaration 
    for (auto& statement : body->statements) {
        if (dynamic_cast<Declaration*>(statement.get())) {
            decl_count++;
        }
        // If
        if (auto* if_stmt = dynamic_cast<If*>(statement.get())) {
            decl_count += count_locals(if_stmt->if_body.get());
            // Else body
            if (if_stmt->else_body.has_value()) {
                 decl_count += count_locals(if_stmt->else_body->get());
            }
        }
        if (auto* while_stmt = dynamic_cast<While*>(statement.get())) {
            decl_count += count_locals(while_stmt->body.get());
        }
        if (auto* body = dynamic_cast<StatementBody*>(statement.get())) {
            decl_count += count_locals(body);
        }
    }
    return decl_count;
}

// Helper function to generate a new label like .L0, .L1, etc.
std::string CodeGenerator::new_label() {
    return ".L" + std::to_string(label_count++) + ":";
}