CodeGen.cpp

// ============================================================
// CodeGen
// ============================================================
//
// Purpose: INTERMEDIATE CODE GENERATION
//          Produces Three-Address Code (TAC) from the AST.
//          Each complex expression is broken into simple steps
//          using temporary variables t0, t1, t2 …
// ============================================================

#include "CodeGen.h"
#include <iostream>
#include <iomanip>
#include <sstream>

// toString() for a single TAC instruction
std::string TACInstruction::toString() const {
    if (op == "GOTO") {
        return "GOTO " + result;               // Unconditional jump
    }
    if (op == "IF_FALSE") {
        return "IF_FALSE " + arg1 + " GOTO " + result;
    }
    if (op == "LABEL") {
        return result + ":";                   // Label definition
    }
    if (op == "PRINT") {
        return "PRINT " + arg1;               // Print instruction
    }
    if (op == "RETURN") {
        return "RETURN " + arg1;
    }
    if (op == "HALT") {
        return "HALT";
    }
    if (arg2.empty()) {
        return result + " = " + arg1;         // Simple assignment/copy
    }
    return result + " = " + arg1 + " " + op + " " + arg2;  // Binary op
}

// CodeGen class implementation
CodeGen::CodeGen() : tempCount_(0), labelCount_(0) {}

std::string CodeGen::newTemp() {
    return "t" + std::to_string(tempCount_++);   // t0, t1, t2, ...
}

std::string CodeGen::newLabel() {
    return "L" + std::to_string(labelCount_++);  // L0, L1, L2, ...
}

// EMIT() – helper to add a new instruction to the list
void CodeGen::emit(const std::string& result,
                   const std::string& op,
                   const std::string& arg1,
                   const std::string& arg2) {
    instructions_.push_back(TACInstruction{result, op, arg1, arg2});
}

// GENERATE() – entry point for code generation, takes the AST and produces TAC
std::vector<TACInstruction>
CodeGen::generate(const std::vector<ASTNodePtr>& stmts) {
    for (const auto& stmt : stmts) {
        genStmt(stmt.get());        // Process each statement
    }
    emit("", "HALT", "", "");       // Program end marker
    return instructions_;
}

// GEN_EXPR() – recursively generate code for an expression node
//    Returns the name of the temp/variable that holds the result
std::string CodeGen::genExpr(const ASTNode* node) {
    // ── Number literal ────────────────────────────────────
    if (auto* n = dynamic_cast<const NumberNode*>(node)) {
        // Strip trailing zeros from integers for cleaner output
        std::string val;
        if (n->value == (int)n->value)
            val = std::to_string((int)n->value);
        else {
            std::ostringstream ss;
            ss << n->value;
            val = ss.str();
        }
        std::string t = newTemp();
        emit(t, "=", val, "");      // t0 = 42
        return t;
    }

    // STRING literal 
    if (auto* n = dynamic_cast<const StringNode*>(node)) {
        std::string t = newTemp();
        emit(t, "=", "\"" + n->value + "\"", "");
        return t;
    }

    // Identifier (variable) – just return its name (assumes it’s already declared)
    if (auto* n = dynamic_cast<const IdentifierNode*>(node)) {
        return n->name;   // Already has a name – use it directly
    }

    // BINARY OPERATOR – generate code for left and right, then combine
    if (auto* n = dynamic_cast<const BinaryOpNode*>(node)) {
        std::string l = genExpr(n->left.get());   // Generate left
        std::string r = genExpr(n->right.get());  // Generate right
        std::string t = newTemp();
        emit(t, n->op, l, r);                     // t2 = t0 + t1
        return t;
    }

    return "";  // Should never reach here for valid AST
}

// GEN_STMT() – recursively generate code for a statement node
void CodeGen::genStmt(const ASTNode* node) {
    // ── Variable declaration ──────────────────────────────
    if (auto* n = dynamic_cast<const VarDeclNode*>(node)) {
        if (n->init) {
            std::string val = genExpr(n->init.get());
            emit(n->name, "=", val, "");   // x = t0
        }
        return;
    }

    // ASSIGNMENT – generate code for the right-hand side, then assign to variable
    if (auto* n = dynamic_cast<const AssignNode*>(node)) {
        std::string val = genExpr(n->expr.get());
        emit(n->name, "=", val, "");        // x = t3
        return;
    }

    // IF statement – generate condition, then branch to then/else blocks
    //  Generated TAC pattern:
    //    <condition>
    //    IF_FALSE cond GOTO elseLabel
    //    <then-body>
    //    GOTO endLabel
    //    elseLabel:
    //    <else-body>
    //    endLabel:
    if (auto* n = dynamic_cast<const IfNode*>(node)) {
        std::string elseLabel = newLabel();
        std::string endLabel  = newLabel();

        std::string cond = genExpr(n->condition.get());
        emit(elseLabel, "IF_FALSE", cond, "");  // Branch if false

        for (const auto& s : n->thenBranch) genStmt(s.get());
        emit(endLabel, "GOTO", "", "");         // Skip else

        emit(elseLabel, "LABEL", "", "");       // Else entry point
        for (const auto& s : n->elseBranch) genStmt(s.get());

        emit(endLabel, "LABEL", "", "");        // Merge point
        return;
    }

    // WHILE loop – generate condition and body with appropriate labels
    //  Generated TAC
    //  Pattern:
    //    loopLabel:
    //    <condition>
    //    IF_FALSE cond GOTO exitLabel
    //    <body>
    //    GOTO loopLabel
    //    exitLabel:
    if (auto* n = dynamic_cast<const WhileNode*>(node)) {
        std::string loopLabel = newLabel();
        std::string exitLabel = newLabel();

        emit(loopLabel, "LABEL", "", "");
        std::string cond = genExpr(n->condition.get());
        emit(exitLabel, "IF_FALSE", cond, "");

        for (const auto& s : n->body) genStmt(s.get());
        emit(loopLabel, "GOTO", "", "");        // Back to loop top

        emit(exitLabel, "LABEL", "", "");
        return;
    }

    // PRINT statement – generate code for the expression, then emit PRINT instruction
    if (auto* n = dynamic_cast<const PrintNode*>(node)) {
        std::string val = genExpr(n->expr.get());
        emit("", "PRINT", val, "");
        return;
    }

    // RETURN statement – generate code for the return value (if any), then emit RETURN instruction
    if (auto* n = dynamic_cast<const ReturnNode*>(node)) {
        if (n->expr) {
            std::string val = genExpr(n->expr.get());
            emit("", "RETURN", val, "");
        } else {
            emit("", "RETURN", "", "");
        }
        return;
    }

    // BLOCK statement – just generate code for each statement in the block
    if (auto* n = dynamic_cast<const BlockNode*>(node)) {
        for (const auto& s : n->statements) genStmt(s.get());
        return;
    }
}

// PRINT_INSTRUCTIONS() – utility to display the generated TAC instructions
void CodeGen::printInstructions() const {
    std::cout << "\n";
    int lineNo = 0;
    for (const auto& instr : instructions_) {
        if (instr.op != "LABEL") {
            // Numbered instruction line
            std::cout << "  " << std::setw(3) << lineNo++ << "  "
                      << instr.toString() << "\n";
        } else {
            // Labels are printed flush-left without a number
            std::cout << "       " << instr.toString() << "\n";
        }
    }
    std::cout << "\n";
}