CodeGenerator.cpp

#include "codegen/CodeGenerator.hpp"
#include "visitor/AstNodeInclude.hpp"

#include <algorithm>
#include <cassert>
#include <cstdarg>
#include <cstdio>

CodeGenerator::CodeGenerator(const std::string &source_file_name,
                             const std::string &save_path,
                             const SymbolManager *const p_symbol_manager)
    : m_symbol_manager_ptr(p_symbol_manager),
      m_source_file_path(source_file_name) {
    // FIXME: assume that the source file is always xxxx.p
    const auto &real_path =
        save_path.empty() ? std::string{"."} : save_path;
    auto slash_pos = source_file_name.rfind("/");
    auto dot_pos = source_file_name.rfind(".");

    if (slash_pos != std::string::npos) {
        ++slash_pos;
    } else {
        slash_pos = 0;
    }
    auto output_file_path{
        real_path + "/" +
        source_file_name.substr(slash_pos, dot_pos - slash_pos) + ".S"};
    m_output_file.reset(fopen(output_file_path.c_str(), "w"));
    assert(m_output_file.get() && "Failed to open output file");
}

static void dumpInstructions(FILE *p_out_file, const char *format, ...) {
    va_list args;
    va_start(args, format);
    vfprintf(p_out_file, format, args);
    va_end(args);
}

void CodeGenerator::visit(ProgramNode &p_program) {
    // Generate RISC-V instructions for program header
    // clang-format off
    constexpr const char *const riscv_assembly_file_prologue =
        "    .file \"%s\"\n"
        "    .option nopic\n";
    // clang-format on
    dumpInstructions(m_output_file.get(), riscv_assembly_file_prologue,
                     m_source_file_path.c_str());

    // Reconstruct the hash table for looking up the symbol entry
    // Hint: Use symbol_manager->lookup(symbol_name) to get the symbol entry.
    m_symbol_manager_ptr->reconstructHashTableFromSymbolTable(
        p_program.getSymbolTable());

    auto visit_ast_node = [&](auto &ast_node) { ast_node->accept(*this); };
    for_each(p_program.getDeclNodes().begin(), p_program.getDeclNodes().end(),
             visit_ast_node);
    for_each(p_program.getFuncNodes().begin(), p_program.getFuncNodes().end(),
             visit_ast_node);

    constexpr const char *const globl_type =
        ".section    .text\n"
        "    .align 2\n"
        "    .globl %s\n"
        "    .type %s, @function\n";
    dumpInstructions(m_output_file.get(), globl_type, "main", "main");
    
    constexpr const char *const main_label = "main:\n";
    dumpInstructions(m_output_file.get(), main_label);
    
    constexpr const char *const prologue =
        "    addi sp, sp, -128\n"
        "    sw ra, 124(sp)\n"
        "    sw s0, 120(sp)\n"
        "    addi s0, sp, 128\n";
    dumpInstructions(m_output_file.get(), prologue);

    m_symbol_manager_ptr->setLocalVariableIndex(0);

    const_cast<CompoundStatementNode &>(p_program.getBody()).accept(*this);

    constexpr const char *const epilogue =
        "    lw ra, 124(sp)\n"
        "    lw s0, 120(sp)\n"
        "    addi sp, sp, 128\n"
        "    jr ra\n"
        "    .size main, .-main\n";
    dumpInstructions(m_output_file.get(), epilogue);

    // Remove the entries in the hash table
    m_symbol_manager_ptr->removeSymbolsFromHashTable(p_program.getSymbolTable());
}

void CodeGenerator::visit(DeclNode &p_decl) {
    p_decl.visitChildNodes(*this);    
}

void CodeGenerator::visit(VariableNode &p_variable) {
    // check if the variable is a function parameter
    if (m_symbol_manager_ptr->lookup(p_variable.getNameCString())->getKind() == SymbolEntry::KindEnum::kParameterKind) {
        // function parameter
        m_symbol_manager_ptr->lookup(p_variable.getNameCString())->setIndex(m_symbol_manager_ptr->getLocalVariableIndex());
        m_symbol_manager_ptr->setLocalVariableIndex(m_symbol_manager_ptr->getLocalVariableIndex() + 1);
        // check if index <= 7
        if (m_symbol_manager_ptr->lookup(p_variable.getNameCString())->getIndex() <= 7) {
            constexpr const char *const function_parameter =
                "    sw a%d, -%d(s0)\n";
            dumpInstructions(m_output_file.get(), function_parameter, m_symbol_manager_ptr->lookup(p_variable.getNameCString())->getIndex(), m_symbol_manager_ptr->lookup(p_variable.getNameCString())->getIndex() * 4 + 12);
        } else {
            constexpr const char *const function_parameter =
                "    lw t0, -%d(s0)\n"
                "    sw t0, -%d(s0)\n";
            dumpInstructions(m_output_file.get(), function_parameter, 128 + (m_symbol_manager_ptr->lookup(p_variable.getNameCString())->getIndex() + 2) * 4, m_symbol_manager_ptr->lookup(p_variable.getNameCString())->getIndex() * 4 + 12);
        }
        return;
    }

    // check if the variable is a global variable
    if (m_symbol_manager_ptr->lookup(p_variable.getNameCString())->getLevel() == 0) {
        // global
        if (m_symbol_manager_ptr->lookup(p_variable.getNameCString())->getKind() == SymbolEntry::KindEnum::kConstantKind) {
            // global constant
            constexpr const char *const global_constant =
                ".section    .rodata\n"
                "    .align 2\n"
                "    .globl %s\n"
                "    .type %s, @object\n"
                "%s:\n"
                "    .word %s\n";
            dumpInstructions(m_output_file.get(), global_constant, p_variable.getNameCString(), p_variable.getNameCString(), p_variable.getNameCString(), p_variable.getConstantPtr()->getConstantValueCString());
        } else {
            // global variable
            constexpr const char *const global_variable =
                ".comm %s, 4, 4\n";
            dumpInstructions(m_output_file.get(), global_variable, p_variable.getNameCString());
        }
    } else {
        // local
        m_symbol_manager_ptr->lookup(p_variable.getNameCString())->setIndex(m_symbol_manager_ptr->getLocalVariableIndex());
        m_symbol_manager_ptr->setLocalVariableIndex(m_symbol_manager_ptr->getLocalVariableIndex() + 1);
        if (m_symbol_manager_ptr->lookup(p_variable.getNameCString())->getKind() == SymbolEntry::KindEnum::kConstantKind) {
            // local constant
            constexpr const char *const local_constant =
                "    addi t0, s0, -%d\n"
                "    addi sp, sp, -4\n"
                "    sw t0, 0(sp)\n"
                "    li t0, %s\n"
                "    addi sp, sp, -4\n"
                "    sw t0, 0(sp)\n"
                "    lw t0, 0(sp)\n"
                "    addi sp, sp, 4\n"
                "    lw t1, 0(sp)\n"
                "    addi sp, sp, 4\n"
                "    sw t0, 0(t1)\n";
            dumpInstructions(m_output_file.get(), local_constant, m_symbol_manager_ptr->lookup(p_variable.getNameCString())->getIndex() * 4 + 12, p_variable.getConstantPtr()->getConstantValueCString());
        } else {
            // local variable
        }
    }
}

void CodeGenerator::visit(ConstantValueNode &p_constant_value) {
    constexpr const char *const constant_value =
        "    li t0, %s\n"
        "    addi sp, sp, -4\n"
        "    sw t0, 0(sp)\n";
    dumpInstructions(m_output_file.get(), constant_value, p_constant_value.getConstantValueCString());
}

void CodeGenerator::visit(FunctionNode &p_function) {
    // Reconstruct the hash table for looking up the symbol entry
    m_symbol_manager_ptr->reconstructHashTableFromSymbolTable(
        p_function.getSymbolTable());

    constexpr const char *const globl_type =
        ".section    .text\n"
        "    .align 2\n"
        "    .globl %s\n"
        "    .type %s, @function\n";
    dumpInstructions(m_output_file.get(), globl_type, p_function.getNameCString(), p_function.getNameCString());

    constexpr const char *const prologue =
        "%s:\n"
        "    addi sp, sp, -128\n"
        "    sw ra, 124(sp)\n"
        "    sw s0, 120(sp)\n"
        "    addi s0, sp, 128\n";
    dumpInstructions(m_output_file.get(), prologue, p_function.getNameCString());

    m_symbol_manager_ptr->setLocalVariableIndex(0);

    p_function.visitChildNodes(*this);
    
    constexpr const char *const epilogue =
        "    lw ra, 124(sp)\n"
        "    lw s0, 120(sp)\n"
        "    addi sp, sp, 128\n"
        "    jr ra\n"
        "    .size %s, .-%s\n";
    dumpInstructions(m_output_file.get(), epilogue, p_function.getNameCString(), p_function.getNameCString());

    // Remove the entries in the hash table
    m_symbol_manager_ptr->removeSymbolsFromHashTable(p_function.getSymbolTable());
}

void CodeGenerator::visit(CompoundStatementNode &p_compound_statement) {
    // Reconstruct the hash table for looking up the symbol entry
    m_symbol_manager_ptr->reconstructHashTableFromSymbolTable(
        p_compound_statement.getSymbolTable());

    if (p_compound_statement.getAttribute() == "ifBody") {
        if (m_symbol_manager_ptr->labelPrinted() == false) {
            constexpr const char *const if_body =
                "L%d:\n";
            dumpInstructions(m_output_file.get(), if_body, m_symbol_manager_ptr->getBranchIndex());
        }
        m_symbol_manager_ptr->setLabelPrinted(false);
        p_compound_statement.visitChildNodes(*this);
    } else if (p_compound_statement.getAttribute() == "elseBody") {
        constexpr const char *const else_body =
            "    j L%d\n"
            "L%d:\n";
        dumpInstructions(m_output_file.get(), else_body, m_symbol_manager_ptr->getBranchIndex() + 2, m_symbol_manager_ptr->getBranchIndex() + 1);
        p_compound_statement.visitChildNodes(*this);
        constexpr const char *const end_if =
            "L%d:\n";
        dumpInstructions(m_output_file.get(), end_if, m_symbol_manager_ptr->getBranchIndex() + 2);
        m_symbol_manager_ptr->setLabelPrinted(true);
        m_symbol_manager_ptr->setBranchIndex(m_symbol_manager_ptr->getBranchIndex() + 3);
    } else if (p_compound_statement.getAttribute() == "whileBody") {
        constexpr const char *const while_body =
            "L%d:\n";
        dumpInstructions(m_output_file.get(), while_body, m_symbol_manager_ptr->getBranchIndex() + 1);
        p_compound_statement.visitChildNodes(*this);
        constexpr const char *const while_body_end =
            "    j L%d\n"
            "L%d:\n";
        dumpInstructions(m_output_file.get(), while_body_end, m_symbol_manager_ptr->getBranchIndex(), m_symbol_manager_ptr->getBranchIndex() + 2);
        m_symbol_manager_ptr->setLabelPrinted(true);
        m_symbol_manager_ptr->setBranchIndex(m_symbol_manager_ptr->getBranchIndex() + 3);
    } else {
        p_compound_statement.visitChildNodes(*this);
    }

    // Remove the entries in the hash table
    m_symbol_manager_ptr->removeSymbolsFromHashTable(
        p_compound_statement.getSymbolTable());
}

void CodeGenerator::visit(PrintNode &p_print) {
    p_print.visitChildNodes(*this);
    constexpr const char *const print =
        "    lw a0, 0(sp)\n"
        "    addi sp, sp, 4\n"
        "    jal ra, printInt\n";
    dumpInstructions(m_output_file.get(), print);
}

void CodeGenerator::visit(BinaryOperatorNode &p_bin_op) {
    // check if operator is plus, minus, multiply or divide
    if (p_bin_op.getOp() == Operator::kPlusOp || p_bin_op.getOp() == Operator::kMinusOp || p_bin_op.getOp() == Operator::kMultiplyOp || p_bin_op.getOp() == Operator::kDivideOp || p_bin_op.getOp() == Operator::kModOp) {
        p_bin_op.visitChildNodes(*this);
        // binary operator
        constexpr const char *const binary_operator =
            "    lw t0, 0(sp)\n"
            "    addi sp, sp, 4\n"
            "    lw t1, 0(sp)\n"
            "    addi sp, sp, 4\n"
            "    %s t0, t1, t0\n"
            "    addi sp, sp, -4\n"
            "    sw t0, 0(sp)\n";
        char op[10];
        switch (p_bin_op.getOp()) {
            case Operator::kPlusOp:
                sprintf(op, "add");
                break;
            case Operator::kMinusOp:
                sprintf(op, "sub");
                break;
            case Operator::kMultiplyOp:
                sprintf(op, "mul");
                break;
            case Operator::kDivideOp:
                sprintf(op, "div");
                break;
            case Operator::kModOp:
                sprintf(op, "rem");
                break;
            default:
                break;
        }
        dumpInstructions(m_output_file.get(), binary_operator, op);
    } else {
        // check if parent is "while"
        if (p_bin_op.getParent() == "while") {
            constexpr const char *const while_condition =
                "L%d:\n";
            m_symbol_manager_ptr->setBranchIndex(m_symbol_manager_ptr->getBranchIndex() - 1);
            if (m_symbol_manager_ptr->labelPrinted() == false) {
                dumpInstructions(m_output_file.get(), while_condition, m_symbol_manager_ptr->getBranchIndex());
            }
        }
        p_bin_op.visitChildNodes(*this);
        // relational operator
        constexpr const char *const relational_operator =
            "    lw t0, 0(sp)\n"
            "    addi sp, sp, 4\n"
            "    lw t1, 0(sp)\n"
            "    addi sp, sp, 4\n"
            "    %s t1, t0, L%d\n";
        char op[10];
        switch (p_bin_op.getOp()) {
            case Operator::kLessOp:
                sprintf(op, "bge");
                break;
            case Operator::kLessOrEqualOp:
                sprintf(op, "bgt");
                break;
            case Operator::kGreaterOp:
                sprintf(op, "ble");
                break;
            case Operator::kGreaterOrEqualOp:
                sprintf(op, "blt");
                break;
            case Operator::kEqualOp:
                sprintf(op, "bne");
                break;
            case Operator::kNotEqualOp:
                sprintf(op, "beq");
                break;
            default:
                break;
        }
        if (p_bin_op.getParent() == "while") {
            dumpInstructions(m_output_file.get(), relational_operator, op, m_symbol_manager_ptr->getBranchIndex() + 2);
        } else {
            dumpInstructions(m_output_file.get(), relational_operator, op, m_symbol_manager_ptr->getBranchIndex() + 1);            
        }
    }
}

void CodeGenerator::visit(UnaryOperatorNode &p_un_op) {
    p_un_op.visitChildNodes(*this);
    constexpr const char *const unary_operator =
        "    lw t0, 0(sp)\n"
        "    addi sp, sp, 4\n"
        "    neg t0, t0\n"
        "    addi sp, sp, -4\n"
        "    sw t0, 0(sp)\n";
    dumpInstructions(m_output_file.get(), unary_operator);
}

void CodeGenerator::visit(FunctionInvocationNode &p_func_invocation) {
    p_func_invocation.visitChildNodes(*this);

    int arg_num = p_func_invocation.getArguments().size();
    for (int i = arg_num - 1; i >= 0; i--) {
        if (i <= 7) {
            constexpr const char *const function_invocation =
                "    lw a%d, 0(sp)\n"
                "    addi sp, sp, 4\n";
            dumpInstructions(m_output_file.get(), function_invocation, i);
        } else {
            constexpr const char *const function_invocation =
                "    lw t0, 0(sp)\n"
                "    addi t1, sp, -%d\n"
                "    sw t0, 0(t1)\n"
                "    addi sp, sp, 4\n";
            dumpInstructions(m_output_file.get(), function_invocation, (i - 7) * 4 + 128);
        }
    }

    constexpr const char *const function_invocation =
        "    jal ra, %s\n"
        "    mv t0, a0\n"
        "    addi sp, sp, -4\n"
        "    sw t0, 0(sp)\n";
    dumpInstructions(m_output_file.get(), function_invocation, p_func_invocation.getNameCString());

}

void CodeGenerator::visit(VariableReferenceNode &p_variable_ref) {
    if (p_variable_ref.getParent() == "print") {
        // global
        if (m_symbol_manager_ptr->lookup(p_variable_ref.getNameCString())->getLevel() == 0) {
            constexpr const char *const global_variable =
                "    la t0, %s\n"
                "    lw t1, 0(t0)\n"
                "    mv t0, t1\n"
                "    addi sp, sp, -4\n"
                "    sw t0, 0(sp)\n";
            dumpInstructions(m_output_file.get(), global_variable, p_variable_ref.getNameCString());
        } else {
            // local
            constexpr const char *const local_variable =
                "    addi t0, s0, -%d\n"
                "    lw t1, 0(t0)\n"
                "    mv t0, t1\n"
                "    addi sp, sp, -4\n"
                "    sw t0, 0(sp)\n";
            dumpInstructions(m_output_file.get(), local_variable, m_symbol_manager_ptr->lookup(p_variable_ref.getNameCString())->getIndex() * 4 + 12);
        }
    } else if (p_variable_ref.getParent() == "binary" || p_variable_ref.getParent() == "FunctionInvocation" || p_variable_ref.getParent() == "return" || p_variable_ref.getParent() == "if" || p_variable_ref.getParent() == "while") {
        if (m_symbol_manager_ptr->lookup(p_variable_ref.getNameCString())->getLevel() == 0) {
            // global
            constexpr const char *const global_variable =
                "    la t0, %s\n"
                "    lw t1, 0(t0)\n"
                "    mv t0, t1\n"
                "    addi sp, sp, -4\n"
                "    sw t0, 0(sp)\n";
            dumpInstructions(m_output_file.get(), global_variable, p_variable_ref.getNameCString());
        } else {
            // local
            constexpr const char *const local_variable =
                "    lw t0, -%d(s0)\n"
                "    addi sp, sp, -4\n"
                "    sw t0, 0(sp)\n";
            dumpInstructions(m_output_file.get(), local_variable, m_symbol_manager_ptr->lookup(p_variable_ref.getNameCString())->getIndex() * 4 + 12);
        }
    } else {
        if (m_symbol_manager_ptr->lookup(p_variable_ref.getNameCString())->getLevel() == 0) {
            // global
            constexpr const char *const global_variable =
                "    la t0, %s\n"
                "    addi sp, sp, -4\n"
                "    sw t0, 0(sp)\n";
            dumpInstructions(m_output_file.get(), global_variable, p_variable_ref.getNameCString());
            p_variable_ref.visitChildNodes(*this);
        } else {
            // local
            constexpr const char *const local_variable =
                "    addi t0, s0, -%d\n"
                "    addi sp, sp, -4\n"
                "    sw t0, 0(sp)\n";
            dumpInstructions(m_output_file.get(), local_variable, m_symbol_manager_ptr->lookup(p_variable_ref.getNameCString())->getIndex() * 4 + 12);
            p_variable_ref.visitChildNodes(*this);
        }
    }
}

void CodeGenerator::visit(AssignmentNode &p_assignment) {
    p_assignment.visitChildNodes(*this);
    constexpr const char *const assignment =
        "    lw t0, 0(sp)\n"
        "    addi sp, sp, 4\n"
        "    lw t1, 0(sp)\n"
        "    addi sp, sp, 4\n"
        "    sw t0, 0(t1)\n";
    dumpInstructions(m_output_file.get(), assignment);
}

void CodeGenerator::visit(ReadNode &p_read) {
    p_read.visitChildNodes(*this);
    constexpr const char *const read =
        "    jal ra, readInt\n"
        "    lw t0, 0(sp)\n"
        "    addi sp, sp, 4\n"
        "    sw a0, 0(t0)\n";
    dumpInstructions(m_output_file.get(), read);
}

void CodeGenerator::visit(IfNode &p_if) {
    p_if.visitChildNodes(*this);
}

void CodeGenerator::visit(WhileNode &p_while) {
    p_while.visitChildNodes(*this);
}

void CodeGenerator::visit(ForNode &p_for) {
    // Reconstruct the hash table for looking up the symbol entry
    m_symbol_manager_ptr->reconstructHashTableFromSymbolTable(
        p_for.getSymbolTable());

    const_cast<DeclNode &>(p_for.getLoopVarDecl()).accept(*this);
    
    const_cast<AssignmentNode &>(p_for.getInitStmt()).accept(*this);

    int branch_index = m_symbol_manager_ptr->getBranchIndex();

    m_symbol_manager_ptr->setBranchIndex(m_symbol_manager_ptr->getBranchIndex() + 3);
    
    constexpr const char *const for_loop_condition =
        "L%d:\n"
        "    lw t0, -%d(s0)\n"
        "    addi sp, sp, -4\n"
        "    sw t0, 0(sp)\n";
    dumpInstructions(m_output_file.get(), for_loop_condition, branch_index, m_symbol_manager_ptr->lookup(p_for.getLoopVarDecl().getFirstId())->getIndex() * 4 + 12);

    const_cast<ExpressionNode &>(p_for.getEndCondition()).accept(*this);

    constexpr const char *const for_loop_condition_end =
        "    lw t0, 0(sp)\n"
        "    addi sp, sp, 4\n"
        "    lw t1, 0(sp)\n"
        "    addi sp, sp, 4\n"
        "    bge t1, t0, L%d\n"
        "L%d:\n";
    dumpInstructions(m_output_file.get(), for_loop_condition_end, branch_index + 2, branch_index + 1);

    const_cast<CompoundStatementNode &>(p_for.getBody()).accept(*this);

    constexpr const char *const for_loop_body_end =
        "    addi t0, s0, -%d\n"
        "    addi sp, sp, -4\n"
        "    sw t0, 0(sp)\n"
        "    lw t0, -%d(s0)\n"
        "    addi sp, sp, -4\n"
        "    sw t0, 0(sp)\n"
        "    li t0, 1\n"
        "    addi sp, sp, -4\n"
        "    sw t0, 0(sp)\n"
        "    lw t0, 0(sp)\n"
        "    addi sp, sp, 4\n"
        "    lw t1, 0(sp)\n"
        "    addi sp, sp, 4\n"
        "    add t0, t1, t0\n"
        "    addi sp, sp, -4\n"
        "    sw t0, 0(sp)\n"
        "    lw t0, 0(sp)\n"
        "    addi sp, sp, 4\n"
        "    lw t1, 0(sp)\n"
        "    addi sp, sp, 4\n"
        "    sw t0, 0(t1)\n"
        "    j L%d\n"
        "L%d:\n";
    dumpInstructions(m_output_file.get(), for_loop_body_end, m_symbol_manager_ptr->lookup(p_for.getLoopVarDecl().getFirstId())->getIndex() * 4 + 12, m_symbol_manager_ptr->lookup(p_for.getLoopVarDecl().getFirstId())->getIndex() * 4 + 12, branch_index, branch_index + 2);


    // Remove the entries in the hash table
    m_symbol_manager_ptr->removeSymbolsFromHashTable(p_for.getSymbolTable());
}

void CodeGenerator::visit(ReturnNode &p_return) {
    p_return.visitChildNodes(*this);
    constexpr const char *const return_node =
        "    lw t0, 0(sp)\n"
        "    addi sp, sp, 4\n"
        "    mv a0, t0\n";
    dumpInstructions(m_output_file.get(), return_node);
}