#include "RISCv64RegAlloc.h"
#include "RISCv64ISel.h"
#include "RISCv64AsmPrinter.h" // For DEBUG output
#include "LegalizeImmediates.h"
#include <algorithm>
#include <vector>
#include <iostream> // For DEBUG output
#include <cassert>  // For assert

namespace sysy {

RISCv64RegAlloc::RISCv64RegAlloc(MachineFunction* mfunc) : MFunc(mfunc) {
    allocable_int_regs = {
        PhysicalReg::T0, PhysicalReg::T1, PhysicalReg::T2, PhysicalReg::T3,
        PhysicalReg::T4, /*PhysicalReg::T5,*/PhysicalReg::T6,
        PhysicalReg::A0, PhysicalReg::A1, PhysicalReg::A2, PhysicalReg::A3,
        PhysicalReg::A4, PhysicalReg::A5, PhysicalReg::A6, PhysicalReg::A7,
        PhysicalReg::S0, PhysicalReg::S1, PhysicalReg::S2, PhysicalReg::S3,
        PhysicalReg::S4, PhysicalReg::S5, PhysicalReg::S6, PhysicalReg::S7,
        PhysicalReg::S8, PhysicalReg::S9, PhysicalReg::S10, PhysicalReg::S11,
    };

    //  创建一个包含所有通用整数寄存器的临时列表
    const std::vector<PhysicalReg> all_int_regs = {
        PhysicalReg::T0, PhysicalReg::T1, PhysicalReg::T2, PhysicalReg::T3,
        PhysicalReg::T4, PhysicalReg::T5, PhysicalReg::T6,
        PhysicalReg::A0, PhysicalReg::A1, PhysicalReg::A2, PhysicalReg::A3,
        PhysicalReg::A4, PhysicalReg::A5, PhysicalReg::A6, PhysicalReg::A7,
        PhysicalReg::S0, PhysicalReg::S1, PhysicalReg::S2, PhysicalReg::S3,
        PhysicalReg::S4, PhysicalReg::S5, PhysicalReg::S6, PhysicalReg::S7,
        PhysicalReg::S8, PhysicalReg::S9, PhysicalReg::S10, PhysicalReg::S11,
    };

    // 映射物理寄存器到特殊的虚拟寄存器ID，用于干扰图中的物理寄存器节点
    // 确保这些特殊ID不会与vreg_counter生成的常规虚拟寄存器ID冲突
    for (PhysicalReg preg : all_int_regs) {
        preg_to_vreg_id_map[preg] = static_cast<unsigned>(PhysicalReg::PHYS_REG_START_ID) + static_cast<unsigned>(preg);
    }
}

// 寄存器分配的主入口点
void RISCv64RegAlloc::run() {
    // 阶段 1: 处理函数调用约定（参数寄存器预着色）
    handleCallingConvention();    
    // 阶段 2: 消除帧索引（为局部变量和栈参数分配栈偏移）
    eliminateFrameIndices();

    // 调试输出当前的LLIR状态
    { // 使用大括号创建一个局部作用域，避免printer变量泄露
        if (DEBUG) {
            std::cerr   << "\n===== LLIR after eliminateFrameIndices for function: " 
                        << MFunc->getName() << " =====\n";
            // 1. 创建一个 AsmPrinter 实例，传入当前的 MachineFunction
            RISCv64AsmPrinter printer(MFunc);
            // 2. 调用 run 方法，将结果打印到标准错误流 (std::cerr)
            // 3. 必须将 debug 参数设为 true！
            //    因为此时指令中仍然包含虚拟寄存器 (%vreg)，
            //    debug模式下的 AsmPrinter 才能正确处理它们而不会报错。
            printer.run(std::cerr, true);
            std::cerr << "===== End of LLIR =====\n\n";
        }
    }

    // 阶段 3: 活跃性分析
    analyzeLiveness();            
    // 阶段 4: 构建干扰图（包含CALL指令对调用者保存寄存器的影响）
    buildInterferenceGraph();     
    // 阶段 5: 图着色算法分配物理寄存器
    colorGraph();                 
    // 阶段 6: 重写函数（插入溢出/填充代码，替换虚拟寄存器为物理寄存器）
    rewriteFunction();            
}

/**
 * @brief 处理调用约定，预先为函数参数和调用返回值分配物理寄存器。
 * 这个函数现在负责处理调用约定的两个方面：
 * 1. 作为被调用者(callee)，如何接收传入的参数。
 * 2. 作为调用者(caller)，如何接收调用的其他函数的返回值。
 */
void RISCv64RegAlloc::handleCallingConvention() {
    Function* F = MFunc->getFunc();
    RISCv64ISel* isel = MFunc->getISel();
    
    // --- 部分1：处理函数传入参数的预着色 ---
    // 获取函数的Argument对象列表
    if (F) {
        auto& args = F->getArguments();
        // RISC-V RV64G调用约定：前8个整型/指针参数通过 a0-a7 传递
        int arg_idx = 0;
        // 遍历 Argument* 列表
        for (Argument* arg : args) {
            if (arg_idx >= 8) {
                break;
            }
            // 获取该 Argument 对象对应的虚拟寄存器ID
            // 通过 MachineFunction -> RISCv64ISel -> vreg_map 来获取
            const auto& vreg_map_from_isel = MFunc->getISel()->getVRegMap();
            assert(vreg_map_from_isel.count(arg) && "Argument not found in ISel's vreg_map!"); 
            // 1. 获取该 Argument 对象对应的虚拟寄存器
            unsigned vreg = isel->getVReg(arg);
            
            // 2. 根据参数索引，确定对应的物理寄存器 (a0, a1, ...)
            auto preg = static_cast<PhysicalReg>(static_cast<int>(PhysicalReg::A0) + arg_idx);
            
            // 3. 在 color_map 中，将 vreg "预着色" 为对应的物理寄存器
            color_map[vreg] = preg;

            arg_idx++;
        }
    }

    // // --- 部分2：[新逻辑] 遍历所有指令，为CALL指令的返回值预着色为 a0 ---
    // // 这是为了强制寄存器分配器知道，call的结果物理上出现在a0寄存器。
    for (auto& mbb : MFunc->getBlocks()) {
        for (auto& instr : mbb->getInstructions()) {
            if (instr->getOpcode() == RVOpcodes::CALL) {
                // 根据协议，如果CALL有返回值，其目标vreg是第一个操作数
                if (!instr->getOperands().empty() && 
                    instr->getOperands().front()->getKind() == MachineOperand::KIND_REG) 
                {
                    auto reg_op = static_cast<RegOperand*>(instr->getOperands().front().get());
                    if (reg_op->isVirtual()) {
                        unsigned ret_vreg = reg_op->getVRegNum();
                        // 强制将这个虚拟寄存器预着色为 a0
                        color_map[ret_vreg] = PhysicalReg::A0;
                        if (DEBUG) {
                            std::cout << "[DEBUG] Pre-coloring vreg" << ret_vreg 
                                      << " to a0 for CALL instruction." << std::endl;
                        }
                    }
                }
            }
        }
    }
}

/**
 * @brief 消除帧索引，为局部变量和栈参数分配栈偏移量，并展开伪指令。
 */
void RISCv64RegAlloc::eliminateFrameIndices() {
    StackFrameInfo& frame_info = MFunc->getFrameInfo();
    // 初始偏移量，为保存ra和s0留出空间。
    // 假设序言是 addi sp, sp, -stack_size; sd ra, stack_size-8(sp); sd s0, stack_size-16(sp);    
    int current_offset = 16;
    
    Function* F = MFunc->getFunc();
    RISCv64ISel* isel = MFunc->getISel();
    
    // 在处理局部变量前，首先为栈参数计算偏移量。
    if (F) {
        int arg_idx = 0;
        for (Argument* arg : F->getArguments()) {
            // 我们只关心第8个索引及之后的参数（即第9个参数开始）
            if (arg_idx >= 8) {
                // 计算偏移量：第一个栈参数(idx=8)在0(s0)，第二个(idx=9)在8(s0)，以此类推。
                int offset = (arg_idx - 8) * 8;
                unsigned vreg = isel->getVReg(arg);

                // 将这个vreg和它的栈偏移存入map。
                // 我们可以复用alloca_offsets，因为它们都代表“vreg到栈偏移”的映射。
                frame_info.alloca_offsets[vreg] = offset;
            }
            arg_idx++;
        }
    }

    // 处理局部变量
    // 遍历AllocaInst来计算局部变量所需的总空间
    for (auto& bb : F->getBasicBlocks()) {
        for (auto& inst : bb->getInstructions()) {
            if (auto alloca = dynamic_cast<AllocaInst*>(inst.get())) {
                // 获取Alloca指令指向的类型 (例如 alloca i32* 中，获取 i32)
                Type* allocated_type = alloca->getType()->as<PointerType>()->getBaseType();
                int size = getTypeSizeInBytes(allocated_type);
                
                // RISC-V要求栈地址8字节对齐
                size = (size + 7) & ~7;
                if (size == 0) size = 8; // 至少分配8字节

                current_offset += size;
                unsigned alloca_vreg = isel->getVReg(alloca);
                // 局部变量使用相对于s0的负向偏移
                frame_info.alloca_offsets[alloca_vreg] = -current_offset;
            }
        }
    }
    frame_info.locals_size = current_offset;

    // 遍历所有机器指令，将伪指令展开为真实指令
    for (auto& mbb : MFunc->getBlocks()) {
        std::vector<std::unique_ptr<MachineInstr>> new_instructions;
        for (auto& instr_ptr : mbb->getInstructions()) {
            RVOpcodes opcode = instr_ptr->getOpcode();

            // --- MODIFICATION START: 处理区分宽度的伪指令 ---
            if (opcode == RVOpcodes::FRAME_LOAD_W || opcode == RVOpcodes::FRAME_LOAD_D) {
                // 确定要生成的真实加载指令是 lw 还是 ld
                RVOpcodes real_load_op = (opcode == RVOpcodes::FRAME_LOAD_W) ? RVOpcodes::LW : RVOpcodes::LD;

                auto& operands = instr_ptr->getOperands();
                unsigned dest_vreg = static_cast<RegOperand*>(operands[0].get())->getVRegNum();
                unsigned alloca_vreg = static_cast<RegOperand*>(operands[1].get())->getVRegNum();
                int offset = frame_info.alloca_offsets.at(alloca_vreg);
                auto addr_vreg = isel->getNewVReg();

                // 展开为: addi addr_vreg, s0, offset
                auto addi = std::make_unique<MachineInstr>(RVOpcodes::ADDI);
                addi->addOperand(std::make_unique<RegOperand>(addr_vreg));
                addi->addOperand(std::make_unique<RegOperand>(PhysicalReg::S0));
                addi->addOperand(std::make_unique<ImmOperand>(offset));
                new_instructions.push_back(std::move(addi));

                // 展开为: lw/ld dest_vreg, 0(addr_vreg)
                auto load_instr = std::make_unique<MachineInstr>(real_load_op);
                load_instr->addOperand(std::make_unique<RegOperand>(dest_vreg));
                load_instr->addOperand(std::make_unique<MemOperand>(
                    std::make_unique<RegOperand>(addr_vreg),
                    std::make_unique<ImmOperand>(0)));
                new_instructions.push_back(std::move(load_instr));

            } else if (opcode == RVOpcodes::FRAME_STORE_W || opcode == RVOpcodes::FRAME_STORE_D) {
                // 确定要生成的真实存储指令是 sw 还是 sd
                RVOpcodes real_store_op = (opcode == RVOpcodes::FRAME_STORE_W) ? RVOpcodes::SW : RVOpcodes::SD;

                auto& operands = instr_ptr->getOperands();
                unsigned src_vreg = static_cast<RegOperand*>(operands[0].get())->getVRegNum();
                unsigned alloca_vreg = static_cast<RegOperand*>(operands[1].get())->getVRegNum();
                int offset = frame_info.alloca_offsets.at(alloca_vreg);
                auto addr_vreg = isel->getNewVReg();
                
                // 展开为: addi addr_vreg, s0, offset
                auto addi = std::make_unique<MachineInstr>(RVOpcodes::ADDI);
                addi->addOperand(std::make_unique<RegOperand>(addr_vreg));
                addi->addOperand(std::make_unique<RegOperand>(PhysicalReg::S0));
                addi->addOperand(std::make_unique<ImmOperand>(offset));
                new_instructions.push_back(std::move(addi));

                // 展开为: sw/sd src_vreg, 0(addr_vreg)
                auto store_instr = std::make_unique<MachineInstr>(real_store_op);
                store_instr->addOperand(std::make_unique<RegOperand>(src_vreg));
                store_instr->addOperand(std::make_unique<MemOperand>(
                    std::make_unique<RegOperand>(addr_vreg),
                    std::make_unique<ImmOperand>(0)));
                new_instructions.push_back(std::move(store_instr));

            } else if (instr_ptr->getOpcode() == RVOpcodes::FRAME_ADDR) { 
                auto& operands = instr_ptr->getOperands();
                unsigned dest_vreg = static_cast<RegOperand*>(operands[0].get())->getVRegNum();
                unsigned alloca_vreg = static_cast<RegOperand*>(operands[1].get())->getVRegNum();
                int offset = frame_info.alloca_offsets.at(alloca_vreg);

                // 将 `frame_addr rd, rs` 展开为 `addi rd, s0, offset`
                auto addi = std::make_unique<MachineInstr>(RVOpcodes::ADDI);
                addi->addOperand(std::make_unique<RegOperand>(dest_vreg));
                addi->addOperand(std::make_unique<RegOperand>(PhysicalReg::S0));
                addi->addOperand(std::make_unique<ImmOperand>(offset));
                new_instructions.push_back(std::move(addi));
            } else {
                new_instructions.push_back(std::move(instr_ptr));
            }
            // --- MODIFICATION END ---
        }
        mbb->getInstructions() = std::move(new_instructions);
    }
}

/**
 * @brief 计算给定 MachineInstr 的 Use (读取) 和 Def (写入) 寄存器集合。
 * 这是活跃性分析的基础。
 * @param instr 要分析的机器指令。
 * @param use 存储 Use 寄存器（虚拟寄存器 ID）的集合。
 * @param def 存储 Def 寄存器（虚拟寄存器 ID）的集合。
 */
void RISCv64RegAlloc::getInstrUseDef(MachineInstr* instr, LiveSet& use, LiveSet& def) {
    bool first_reg_is_def = true; // 默认情况下，指令的第一个寄存器操作数是定义 (def)
    auto opcode = instr->getOpcode();

    // 1. 特殊指令的 `is_def` 标志调整
    // 这些指令的第一个寄存器操作数是源操作数 (use)，而不是目标操作数 (def)。
    if (opcode == RVOpcodes::SW || opcode == RVOpcodes::SD || 
        opcode == RVOpcodes::BEQ || opcode == RVOpcodes::BNE || 
        opcode == RVOpcodes::BLT || opcode == RVOpcodes::BGE ||
        opcode == RVOpcodes::BLTU || opcode == RVOpcodes::BGEU ||
        opcode == RVOpcodes::RET || opcode == RVOpcodes::J) {
        first_reg_is_def = false;
    }

    // JAL 和 JALR 指令定义 ra (x1)
    if (opcode == RVOpcodes::JAL || opcode == RVOpcodes::JALR) {
        // 使用 ra 对应的特殊虚拟寄存器ID
        def.insert(preg_to_vreg_id_map.at(PhysicalReg::RA)); 
        first_reg_is_def = false; // JAL/JALR 的第一个操作数是 ra，已经处理为 def
    }
    
    // 2. CALL 指令的特殊处理
    if (opcode == RVOpcodes::CALL) {
        // 根据 s1 分支 ISel 定义的协议来解析操作数列表
        bool first_reg_operand_is_def = true;
        for (auto& op : instr->getOperands()) {
            if (op->getKind() == MachineOperand::KIND_REG) {
                auto reg_op = static_cast<RegOperand*>(op.get());
                if (reg_op->isVirtual()) {
                    // 协议：第一个寄存器操作数是返回值 (def)
                    if (first_reg_operand_is_def) {
                        def.insert(reg_op->getVRegNum());
                        first_reg_operand_is_def = false;
                    } else {
                        // 后续所有寄存器操作数都是参数 (use)
                        use.insert(reg_op->getVRegNum());
                    }
                } else { // [修复] CALL指令也可能定义物理寄存器（如a0）
                    if (first_reg_operand_is_def) {
                         if (preg_to_vreg_id_map.count(reg_op->getPReg())) {
                            def.insert(preg_to_vreg_id_map.at(reg_op->getPReg()));
                        }
                        first_reg_operand_is_def = false;
                    } else {
                         if (preg_to_vreg_id_map.count(reg_op->getPReg())) {
                            use.insert(preg_to_vreg_id_map.at(reg_op->getPReg()));
                        }
                    }
                }
            }
        }
        return; // CALL 指令处理完毕
    }
    
    // 3. 对其他所有指令的通用处理逻辑 [已重构和修复]
    for (const auto& op : instr->getOperands()) {
        if (op->getKind() == MachineOperand::KIND_REG) {
            auto reg_op = static_cast<RegOperand*>(op.get());
            
            if (first_reg_is_def) {
                // --- 处理定义（Def） ---
                if (reg_op->isVirtual()) {
                    def.insert(reg_op->getVRegNum());
                } else { // 物理寄存器也可以是 Def
                    if (preg_to_vreg_id_map.count(reg_op->getPReg())) {
                        def.insert(preg_to_vreg_id_map.at(reg_op->getPReg()));
                    }
                }
                first_reg_is_def = false; // **关键**：处理完第一个寄存器后，立即更新标志
            } else {
                // --- 处理使用（Use） ---
                if (reg_op->isVirtual()) {
                    use.insert(reg_op->getVRegNum());
                } else { // 物理寄存器也可以是 Use
                    if (preg_to_vreg_id_map.count(reg_op->getPReg())) {
                        use.insert(preg_to_vreg_id_map.at(reg_op->getPReg()));
                    }
                }
            }
        } else if (op->getKind() == MachineOperand::KIND_MEM) {
            // [保持不变] 内存操作数的处理逻辑看起来是正确的
            auto mem_op = static_cast<MemOperand*>(op.get());
            auto base_reg = mem_op->getBase();
            if (base_reg->isVirtual()) {
                use.insert(base_reg->getVRegNum());
            } else {
                PhysicalReg preg = base_reg->getPReg();
                if (preg_to_vreg_id_map.count(preg)) {
                    use.insert(preg_to_vreg_id_map.at(preg));
                }
            }
            
            // 对于存储内存指令 (SW, SD)，要存储的值（第一个操作数）也是 `use`
            if ((opcode == RVOpcodes::SW || opcode == RVOpcodes::SD) && 
                 !instr->getOperands().empty() &&
                 instr->getOperands().front()->getKind() == MachineOperand::KIND_REG) {
                auto src_reg_op = static_cast<RegOperand*>(instr->getOperands().front().get());
                if (src_reg_op->isVirtual()) {
                    use.insert(src_reg_op->getVRegNum());
                } else {
                    if (preg_to_vreg_id_map.count(src_reg_op->getPReg())) {
                       use.insert(preg_to_vreg_id_map.at(src_reg_op->getPReg()));
                    }
                }
            }
        }
    }
}

/**
 * @brief 计算一个类型在内存中占用的字节数。
 * @param type 需要计算大小的IR类型。
 * @return 该类型占用的字节数。
 */
unsigned RISCv64RegAlloc::getTypeSizeInBytes(Type* type) {
    if (!type) {
        assert(false && "Cannot get size of a null type.");
        return 0;
    }

    switch (type->getKind()) {
        // 对于SysY语言，基本类型int和float都占用4字节
        case Type::kInt:
        case Type::kFloat:
            return 4;

        // 指针类型在RISC-V 64位架构下占用8字节
        // 虽然SysY没有'int*'语法，但数组变量在IR层面本身就是指针类型
        case Type::kPointer:
            return 8;

        // 数组类型的总大小 = 元素数量 * 单个元素的大小
        case Type::kArray: {
            auto arrayType = type->as<ArrayType>();
            // 递归调用以计算元素大小
            return arrayType->getNumElements() * getTypeSizeInBytes(arrayType->getElementType());
        }

        // 其他类型，如Void, Label等不占用栈空间，或者不应该出现在这里
        default:
            // 如果遇到未处理的类型，触发断言，方便调试
            assert(false && "Unsupported type for size calculation.");
            return 0;
    }
}

void RISCv64RegAlloc::analyzeLiveness() {
    // === 阶段 1: 预计算每个基本块的 use 和 def 集合 ===
    // 这样可以避免在主循环中重复计算
    std::map<MachineBasicBlock*, LiveSet> block_uses;
    std::map<MachineBasicBlock*, LiveSet> block_defs;
    for (auto& mbb_ptr : MFunc->getBlocks()) {
        MachineBasicBlock* mbb = mbb_ptr.get();
        LiveSet uses, defs;
        for (auto& instr_ptr : mbb->getInstructions()) {
            LiveSet instr_use, instr_def;
            getInstrUseDef(instr_ptr.get(), instr_use, instr_def);
            // use[B] = use[B] U (instr_use - def[B])
            for (unsigned u : instr_use) {
                if (defs.find(u) == defs.end()) {
                    uses.insert(u);
                }
            }
            // def[B] = def[B] U instr_def
            defs.insert(instr_def.begin(), instr_def.end());
        }
        block_uses[mbb] = uses;
        block_defs[mbb] = defs;
    }

    // === 阶段 2: 在“块”粒度上进行迭代数据流分析，直到收敛 ===
    std::map<MachineBasicBlock*, LiveSet> block_live_in;
    std::map<MachineBasicBlock*, LiveSet> block_live_out;
    bool changed = true;
    while (changed) {
        changed = false;
        // 以逆后序遍历基本块，可以加速收敛，但简单的逆序对于大多数情况也有效
        for (auto it = MFunc->getBlocks().rbegin(); it != MFunc->getBlocks().rend(); ++it) {
            auto& mbb = *it;
            
            // 2.1 计算 live_out[B] = U_{S in succ(B)} live_in[S]
            LiveSet new_live_out;
            for (auto succ : mbb->successors) {
                new_live_out.insert(block_live_in[succ].begin(), block_live_in[succ].end());
            }

            // 2.2 计算 live_in[B] = use[B] U (live_out[B] - def[B])
            LiveSet live_out_minus_def = new_live_out;
            for (unsigned d : block_defs[mbb.get()]) {
                live_out_minus_def.erase(d);
            }
            LiveSet new_live_in = block_uses[mbb.get()];
            new_live_in.insert(live_out_minus_def.begin(), live_out_minus_def.end());

            // 2.3 检查 live_in 和 live_out 是否变化，以判断是否达到不动点
            if (block_live_out[mbb.get()] != new_live_out) {
                changed = true;
                block_live_out[mbb.get()] = new_live_out;
            }
            if (block_live_in[mbb.get()] != new_live_in) {
                changed = true;
                block_live_in[mbb.get()] = new_live_in;
            }
        }
    }

    // === 阶段 3: 进行一次指令粒度的遍历，填充最终的 live_in_map 和 live_out_map ===
    // 此时块级别的活跃信息已经稳定，我们只需遍历一次即可
    for (auto& mbb_ptr : MFunc->getBlocks()) {
        MachineBasicBlock* mbb = mbb_ptr.get();
        LiveSet live_out = block_live_out[mbb]; // 从已收敛的块级 live_out 开始

        for (auto instr_it = mbb->getInstructions().rbegin(); instr_it != mbb->getInstructions().rend(); ++instr_it) {
            MachineInstr* instr = instr_it->get();
            live_out_map[instr] = live_out;

            LiveSet use, def;
            getInstrUseDef(instr, use, def);

            LiveSet live_in = use;
            LiveSet diff = live_out;
            for (auto vreg : def) {
                diff.erase(vreg);
            }
            live_in.insert(diff.begin(), diff.end());
            live_in_map[instr] = live_in;
            
            // 更新 live_out，为块内的上一条指令做准备
            live_out = live_in;
        }
    }
}

// 辅助函数，用于清晰地打印寄存器集合。可以放在 .cpp 文件的顶部。
void RISCv64RegAlloc::printLiveSet(const LiveSet& s, const std::string& name, std::ostream& os) {
    os << "    " << name << ": { ";
    for (unsigned vreg : s) {
        // 为了可读性，将物理寄存器对应的特殊ID进行转换
        if (vreg >= static_cast<unsigned>(sysy::PhysicalReg::PHYS_REG_START_ID)) {
            os << "preg(" << (vreg - static_cast<unsigned>(sysy::PhysicalReg::PHYS_REG_START_ID)) << ") ";
        } else {
            os << "%vreg" << vreg << " ";
        }
    }
    os << "}\n";
}

void RISCv64RegAlloc::buildInterferenceGraph() {
    std::set<unsigned> all_vregs;
    // 收集所有虚拟寄存器和物理寄存器在干扰图中的节点ID
    for (auto& mbb : MFunc->getBlocks()) {
        for(auto& instr : mbb->getInstructions()) {
            LiveSet use, def;
            getInstrUseDef(instr.get(), use, def);
            for(auto u : use) all_vregs.insert(u);
            for(auto d : def) all_vregs.insert(d);
        }
    }
    // 添加所有物理寄存器对应的特殊虚拟寄存器ID到all_vregs，作为干扰图节点
    for (auto preg : allocable_int_regs) {
        all_vregs.insert(preg_to_vreg_id_map.at(preg));
    }

    // 初始化干扰图邻接表
    for (auto vreg : all_vregs) { interference_graph[vreg] = {}; }

    // 创建一个临时的AsmPrinter用于打印指令，方便调试
    RISCv64AsmPrinter temp_printer(MFunc);
    temp_printer.setStream(std::cerr);

    for (auto& mbb : MFunc->getBlocks()) {
        if (DEEPDEBUG) std::cerr << "--- Building Graph for Basic Block: " << mbb->getName() << " ---\n";
        for (auto& instr_ptr : mbb->getInstructions()) {
            MachineInstr* instr = instr_ptr.get();
            if (DEEPDEBUG) {
                // 打印当前正在处理的指令
                std::cerr << "  Instr: ";
                temp_printer.printInstruction(instr, true); // 使用 true 来打印虚拟寄存器
            }
            
            LiveSet def, use;
            getInstrUseDef(instr, use, def);
            const LiveSet& live_out = live_out_map.at(instr);

            // [新增调试逻辑] 打印所有相关的寄存器集合
            if (DEEPDEBUG) {
                printLiveSet(use, "Use     ", std::cerr);
                printLiveSet(def, "Def     ", std::cerr);
                printLiveSet(live_out, "Live_Out", std::cerr); // 这是我们最关心的信息
            }

            // 标准干扰图构建：def 与 live_out 中的其他变量干扰
            for (unsigned d : def) {
                for (unsigned l : live_out) {
                    if (d != l) {
                        // [新增调试逻辑] 打印添加的干扰边及其原因
                        if (DEEPDEBUG && interference_graph[d].find(l) == interference_graph[d].end()) {
                           std::cerr << "    Edge (Def-LiveOut): %vreg" << d << " <-> %vreg" << l << "\n";
                        }
                        interference_graph[d].insert(l);
                        interference_graph[l].insert(d);
                    }
                }
            }

            // 在非move指令中，def 与 use 互相干扰
            if (instr->getOpcode() != RVOpcodes::MV) {
                for (unsigned d : def) {
                    for (unsigned u : use) {
                        if (d != u) {
                            // [新增调试逻辑] 打印添加的干扰边及其原因
                            if (DEEPDEBUG && interference_graph[d].find(u) == interference_graph[d].end()) {
                                std::cerr << "    Edge (Def-Use)    : %vreg" << d << " <-> %vreg" << u << "\n";
                            }
                            interference_graph[d].insert(u);
                            interference_graph[u].insert(d);
                        }
                    }
                }
            }

            // *** 处理 CALL 指令的隐式 def ***
            if (instr->getOpcode() == RVOpcodes::CALL) {
                // 你的原始CALL调试信息
                if (DEEPDEBUG) {
                    std::string live_out_str;
                    for (unsigned vreg : live_out) {
                        live_out_str += "%vreg" + std::to_string(vreg) + " ";
                    }
                    std::cerr << "[DEEPDEBUG] buildInterferenceGraph: CALL instruction found. Live out set is: {" 
                              << live_out_str << "}" << std::endl;
                }
                // CALL 指令会定义（杀死）所有调用者保存的寄存器。
                // 因此，所有调用者保存的物理寄存器都与 CALL 指令的 live_out 中的所有变量冲突。
                const std::vector<PhysicalReg>& caller_saved_regs = getCallerSavedIntRegs();
                for (PhysicalReg cs_reg : caller_saved_regs) {
                    if (preg_to_vreg_id_map.count(cs_reg)) {
                        unsigned cs_vreg_id = preg_to_vreg_id_map.at(cs_reg); // 获取物理寄存器对应的特殊vreg ID
                    
                        // 将这个物理寄存器节点与 CALL 指令的 live_out 中的所有虚拟寄存器添加干扰边。
                        for (unsigned live_vreg_out : live_out) {
                            if (cs_vreg_id != live_vreg_out) { // 避免自己和自己干扰
                                // [新增调试逻辑] 打印添加的干扰边及其原因
                                if (DEEPDEBUG && interference_graph[cs_vreg_id].find(live_vreg_out) == interference_graph[cs_vreg_id].end()) {
                                    std::cerr << "    Edge (CALL)       : preg(" << static_cast<int>(cs_reg) << ") <-> %vreg" << live_vreg_out << "\n";
                                }
                                interference_graph[cs_vreg_id].insert(live_vreg_out);
                                interference_graph[live_vreg_out].insert(cs_vreg_id);
                            }
                        }
                    } else {
                        // 如果物理寄存器没有对应的特殊虚拟寄存器ID，可能是因为它不是调用者保存的寄存器。
                        // 这种情况通常不应该发生，但我们可以在这里添加一个警告或错误处理。
                        if (DEEPDEBUG) {
                            std::cerr << "Warning: Physical register " << static_cast<int>(cs_reg) 
                                      << " does not have a corresponding special vreg ID.\n";
                        }
                    }
                }
            }
            if (DEEPDEBUG) std::cerr << "  ----------------\n";
        }
    }
}

void RISCv64RegAlloc::colorGraph() {
    std::vector<unsigned> sorted_vregs;
    for (auto const& [vreg, neighbors] : interference_graph) {
        // 只为未预着色的虚拟寄存器排序和着色
        if (color_map.find(vreg) == color_map.end() && vreg < static_cast<unsigned>(PhysicalReg::PHYS_REG_START_ID)) {
            sorted_vregs.push_back(vreg);
        }
    }

    // 排序
    std::sort(sorted_vregs.begin(), sorted_vregs.end(), [&](unsigned a, unsigned b) {
        return interference_graph[a].size() > interference_graph[b].size();
    });

    // 着色
    for (unsigned vreg : sorted_vregs) {
        std::set<PhysicalReg> used_colors;
        for (unsigned neighbor_id : interference_graph.at(vreg)) {
            // --- 关键改进 (来自 rec 分支) ---
            
            // 情况 1: 邻居是一个已经被着色的虚拟寄存器
            if (color_map.count(neighbor_id)) {
                used_colors.insert(color_map.at(neighbor_id));
            } 
            // 情况 2: 邻居本身就是一个代表物理寄存器的节点
            else if (neighbor_id >= static_cast<unsigned>(PhysicalReg::PHYS_REG_START_ID)) {
                // 从特殊ID反向解析出是哪个物理寄存器
                PhysicalReg neighbor_preg = static_cast<PhysicalReg>(neighbor_id - static_cast<unsigned>(PhysicalReg::PHYS_REG_START_ID));
                used_colors.insert(neighbor_preg);
            }
        }
        
        bool colored = false;
        for (PhysicalReg preg : allocable_int_regs) {
            if (used_colors.find(preg) == used_colors.end()) {
                color_map[vreg] = preg;
                colored = true;
                break;
            }
        }
        if (!colored) {
            spilled_vregs.insert(vreg);
        }
    }
}

void RISCv64RegAlloc::rewriteFunction() {
    StackFrameInfo& frame_info = MFunc->getFrameInfo();
    int current_offset = frame_info.locals_size;

    // --- FIX 1: 动态计算溢出槽大小 ---
    // 根据溢出虚拟寄存器的真实类型，为其在栈上分配正确大小的空间。
    for (unsigned vreg : spilled_vregs) {
        // 从反向映射中查找 vreg 对应的 IR Value
        assert(vreg_to_value_map.count(vreg) && "Spilled vreg not found in map!");
        Value* val = vreg_to_value_map.at(vreg);
        
        // 使用辅助函数获取类型大小
        int size = getTypeSizeInBytes(val->getType());
        
        // 保持栈8字节对齐
        current_offset += size;
        current_offset = (current_offset + 7) & ~7;

        frame_info.spill_offsets[vreg] = -current_offset;
    }
    frame_info.spill_size = current_offset - frame_info.locals_size;

    for (auto& mbb : MFunc->getBlocks()) {
        std::vector<std::unique_ptr<MachineInstr>> new_instructions;
        for (auto& instr_ptr : mbb->getInstructions()) {
            LiveSet use, def;
            getInstrUseDef(instr_ptr.get(), use, def);

            // --- FIX 2: 为溢出的 'use' 操作数插入正确的加载指令 ---
            for (unsigned vreg : use) {
                if (spilled_vregs.count(vreg)) {
                    // 同样地，根据 vreg 的类型决定使用 lw 还是 ld
                    assert(vreg_to_value_map.count(vreg));
                    Value* val = vreg_to_value_map.at(vreg);
                    RVOpcodes load_op = val->getType()->isPointer() ? RVOpcodes::LD : RVOpcodes::LW;

                    int offset = frame_info.spill_offsets.at(vreg);
                    auto load = std::make_unique<MachineInstr>(load_op);
                    load->addOperand(std::make_unique<RegOperand>(vreg));
                    load->addOperand(std::make_unique<MemOperand>(
                        std::make_unique<RegOperand>(PhysicalReg::S0),
                        std::make_unique<ImmOperand>(offset)
                    ));
                    new_instructions.push_back(std::move(load));
                }
            }

            new_instructions.push_back(std::move(instr_ptr));

            // --- FIX 3: 为溢出的 'def' 操作数插入正确的存储指令 ---
            for (unsigned vreg : def) {
                if (spilled_vregs.count(vreg)) {
                    // 根据 vreg 的类型决定使用 sw 还是 sd
                    assert(vreg_to_value_map.count(vreg));
                    Value* val = vreg_to_value_map.at(vreg);
                    RVOpcodes store_op = val->getType()->isPointer() ? RVOpcodes::SD : RVOpcodes::SW;

                    int offset = frame_info.spill_offsets.at(vreg);
                    auto store = std::make_unique<MachineInstr>(store_op);
                    store->addOperand(std::make_unique<RegOperand>(vreg));
                    store->addOperand(std::make_unique<MemOperand>(
                        std::make_unique<RegOperand>(PhysicalReg::S0),
                        std::make_unique<ImmOperand>(offset)
                    ));
                    new_instructions.push_back(std::move(store));
                }
            }
        }
        mbb->getInstructions() = std::move(new_instructions);
    }
    
    // 最后的虚拟寄存器到物理寄存器的替换过程保持不变
    for (auto& mbb : MFunc->getBlocks()) {
        for (auto& instr_ptr : mbb->getInstructions()) {
            for (auto& op_ptr : instr_ptr->getOperands()) {

                // 情况一：操作数本身就是一个寄存器 (例如 add rd, rs1, rs2 中的所有操作数)
                if(op_ptr->getKind() == MachineOperand::KIND_REG) {
                    auto reg_op = static_cast<RegOperand*>(op_ptr.get());
                    if (reg_op->isVirtual()) {
                        unsigned vreg = reg_op->getVRegNum();
                        if (color_map.count(vreg)) {
                            PhysicalReg preg = color_map.at(vreg);
                            reg_op->setPReg(preg);
                        } else if (spilled_vregs.count(vreg)) {
                            // 如果vreg被溢出，替换为专用的溢出物理寄存器t6
                            reg_op->setPReg(PhysicalReg::T6);
                        }
                    }
                } 
                // 情况二：操作数是一个内存地址 (例如 lw rd, offset(rs1) 中的 offset(rs1))
                else if (op_ptr->getKind() == MachineOperand::KIND_MEM) {
                    auto mem_op = static_cast<MemOperand*>(op_ptr.get());
                    // 获取内存操作数内部的“基址寄存器”
                    auto base_reg_op = mem_op->getBase();

                    // 对这个基址寄存器，执行与情况一完全相同的替换逻辑
                    if(base_reg_op->isVirtual()){
                        unsigned vreg = base_reg_op->getVRegNum();
                        if(color_map.count(vreg)) {
                            // 如果基址vreg被成功着色，替换
                            PhysicalReg preg = color_map.at(vreg);
                            base_reg_op->setPReg(preg);

                        } else if (spilled_vregs.count(vreg)) {
                            // 如果基址vreg被溢出，替换为t6
                            base_reg_op->setPReg(PhysicalReg::T6);
                        }
                    }
                }
            }
        }
    }
}

} // namespace sysy