// in file: RISCv64RegAlloc.cpp

#include "RISCv64RegAlloc.h"
#include "RISCv64AsmPrinter.h"
#include <algorithm>
#include <iostream>
#include <sstream>
#include <cassert>

namespace sysy {

// 构造函数：初始化寄存器池和数据结构
RISCv64RegAlloc::RISCv64RegAlloc(MachineFunction* mfunc)
    : MFunc(mfunc), 
      ISel(mfunc->getISel()),
      vreg_to_value_map(ISel->getVRegValueMap()),
      vreg_type_map(ISel->getVRegTypeMap()) {
    // 1. 初始化可分配的整数寄存器池
    allocable_int_regs = {
        PhysicalReg::T0, PhysicalReg::T1, PhysicalReg::T2, PhysicalReg::T3, PhysicalReg::T4, /* T5保留 */ PhysicalReg::T6,
        PhysicalReg::A0, PhysicalReg::A1, PhysicalReg::A2, PhysicalReg::A3, PhysicalReg::A4, PhysicalReg::A5, PhysicalReg::A6, PhysicalReg::A7,
        PhysicalReg::S1, PhysicalReg::S2, PhysicalReg::S3, PhysicalReg::S4, PhysicalReg::S5, PhysicalReg::S6, PhysicalReg::S7,
        PhysicalReg::S8, PhysicalReg::S9, PhysicalReg::S10, PhysicalReg::S11,
        // S0 是帧指针，不参与分配
    };
    K_int = allocable_int_regs.size();

    // 2. 初始化可分配的浮点寄存器池
    allocable_fp_regs = {
        PhysicalReg::F0, PhysicalReg::F1, PhysicalReg::F2, PhysicalReg::F3, PhysicalReg::F4, PhysicalReg::F5, PhysicalReg::F6, PhysicalReg::F7,
        PhysicalReg::F10, PhysicalReg::F11, PhysicalReg::F12, PhysicalReg::F13, PhysicalReg::F14, PhysicalReg::F15, PhysicalReg::F16, PhysicalReg::F17,
        PhysicalReg::F8, PhysicalReg::F9, PhysicalReg::F18, PhysicalReg::F19, PhysicalReg::F20, PhysicalReg::F21, PhysicalReg::F22,
        PhysicalReg::F23, PhysicalReg::F24, PhysicalReg::F25, PhysicalReg::F26, PhysicalReg::F27,
        PhysicalReg::F28, PhysicalReg::F29, PhysicalReg::F30, PhysicalReg::F31,
    };
    K_fp = allocable_fp_regs.size();

    // 3. 预着色所有物理寄存器
    const unsigned offset = static_cast<unsigned>(PhysicalReg::PHYS_REG_START_ID);
    for (const auto& reg : allocable_int_regs) precolored.insert(offset + static_cast<unsigned>(reg));
    for (const auto& reg : allocable_fp_regs) precolored.insert(offset + static_cast<unsigned>(reg));
    precolored.insert(offset + static_cast<unsigned>(PhysicalReg::S0));
    precolored.insert(offset + static_cast<unsigned>(PhysicalReg::RA));
    precolored.insert(offset + static_cast<unsigned>(PhysicalReg::SP));
    precolored.insert(offset + static_cast<unsigned>(PhysicalReg::ZERO));
}

// 主入口: 迭代运行分配算法直到无溢出
void RISCv64RegAlloc::run() {
    if (DEBUG) std::cerr << "===== LLIR Before Running Graph Coloring Register Allocation " << MFunc->getName() << " =====\n";
    std::stringstream ss_before_reg_alloc;
    if (DEBUG) {
        RISCv64AsmPrinter printer_reg_alloc(MFunc);
        printer_reg_alloc.run(ss_before_reg_alloc, true);
        std::cout << ss_before_reg_alloc.str();
    }

    if (DEBUG) std::cerr << "===== Running Graph Coloring Register Allocation for function: " << MFunc->getName() << " =====\n";
    
    while (true) {
        if (doAllocation()) {
            break;
        } else {
            rewriteProgram();
            if (DEBUG) std::cerr << "--- Spilling detected, re-running allocation ---\n";
        }
    }

    applyColoring();
    
    MFunc->getFrameInfo().vreg_to_preg_map = this->color_map;
    collectUsedCalleeSavedRegs();
    if (DEBUG) std::cerr << "===== Finished Graph Coloring Register Allocation =====\n\n";
}

// 单次分配的核心流程
bool RISCv64RegAlloc::doAllocation() {
    initialize();
    precolorByCallingConvention();
    analyzeLiveness();
    build();
    makeWorklist();

    while (!simplifyWorklist.empty() || !worklistMoves.empty() || !freezeWorklist.empty() || !spillWorklist.empty()) {
        if (DEEPDEBUG) dumpState("Loop Start");
        if (!simplifyWorklist.empty()) simplify();
        else if (!worklistMoves.empty()) coalesce();
        else if (!freezeWorklist.empty()) freeze();
        else if (!spillWorklist.empty()) selectSpill();
    }

    if (DEEPDEBUG) dumpState("Before AssignColors");
    assignColors();
    return spilledNodes.empty();
}

void RISCv64RegAlloc::precolorByCallingConvention() {
    // [新增] 在处理前，先清空颜色相关的状态，确保重试时不会出错
    color_map.clear();
    coloredNodes.clear();

    Function* F = MFunc->getFunc();
    if (!F) return;

    // --- 部分1：处理函数传入参数的预着色 ---
    int int_arg_idx = 0;
    int float_arg_idx = 0;

    for (Argument* arg : F->getArguments()) {
        unsigned vreg = ISel->getVReg(arg);
        
        if (arg->getType()->isFloat()) {
            if (float_arg_idx < 8) { // fa0-fa7
                auto preg = static_cast<PhysicalReg>(static_cast<int>(PhysicalReg::F10) + float_arg_idx);
                color_map[vreg] = preg;
                float_arg_idx++;
            }
        } else { // 整数或指针
            if (int_arg_idx < 8) { // a0-a7
                auto preg = static_cast<PhysicalReg>(static_cast<int>(PhysicalReg::A0) + int_arg_idx);
                color_map[vreg] = preg;
                int_arg_idx++;
            }
        }
    }

    // --- 部分2：为CALL指令的返回值预着色 ---
    for (auto& mbb : MFunc->getBlocks()) {
        for (auto& instr : mbb->getInstructions()) {
            if (instr->getOpcode() == RVOpcodes::CALL) {
                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();
                        assert(vreg_to_value_map.count(ret_vreg) && "Return vreg not found!");
                        Value* ret_val = vreg_to_value_map.at(ret_vreg);

                        if (ret_val->getType()->isFloat()) {
                            color_map[ret_vreg] = PhysicalReg::F10; // fa0
                        } else {
                            color_map[ret_vreg] = PhysicalReg::A0;  // a0
                        }
                    }
                }
            }
        }
    }
    
    // 将所有预着色的vreg视为已着色节点
    for(const auto& pair : color_map) {
        coloredNodes.insert(pair.first);
    }
    if (DEEPDEBUG) {
        std::cerr << "Precolored registers: { ";
        // 修改部分：将物理寄存器ID转换为其字符串名称
        for (unsigned v : precolored) std::cerr << regIdToString(v) << " ";
        std::cerr << "}\nColored nodes: { ";
        for (unsigned v : coloredNodes) std::cerr << "%vreg" << v << " ";
        std::cerr << "}\n";
    }
}

// 初始化/重置所有数据结构
void RISCv64RegAlloc::initialize() {
    initial.clear();
    simplifyWorklist.clear();
    freezeWorklist.clear();
    spillWorklist.clear();
    spilledNodes.clear();
    coalescedNodes.clear();
    coloredNodes.clear(); // [修正] 恢复对这两个集合的清除
    selectStack.clear();

    coalescedMoves.clear();
    constrainedMoves.clear();
    frozenMoves.clear();
    worklistMoves.clear();
    activeMoves.clear();

    adjList.clear();
    degree.clear();
    moveList.clear();
    alias.clear();
    color_map.clear(); // [修正] 恢复对 color_map 的清除
}

// 活跃性分析（此部分为标准数据流分析，与现有版本类似但更精细）
void RISCv64RegAlloc::analyzeLiveness() {
    live_in_map.clear();
    live_out_map.clear();
    
    std::map<const MachineBasicBlock*, VRegSet> block_uses;
    std::map<const MachineBasicBlock*, VRegSet> block_defs;

    for (const auto& mbb_ptr : MFunc->getBlocks()) {
        const MachineBasicBlock* mbb = mbb_ptr.get();
        VRegSet uses, defs;
        for (const auto& instr_ptr : mbb->getInstructions()) {
            VRegSet instr_use, instr_def;
            // 使用新的、能看到物理寄存器的版本
            getInstrUseDef_Liveness(instr_ptr.get(), instr_use, instr_def);
            for (unsigned u : instr_use) {
                if (defs.find(u) == defs.end()) uses.insert(u);
            }
            defs.insert(instr_def.begin(), instr_def.end());
        }
        block_uses[mbb] = uses;
        block_defs[mbb] = defs;
    }

    bool changed = true;
    std::map<const MachineBasicBlock*, VRegSet> live_in, live_out;
    while (changed) {
        changed = false;
        for (auto it = MFunc->getBlocks().rbegin(); it != MFunc->getBlocks().rend(); ++it) {
            const auto& mbb_ptr = *it;
            const MachineBasicBlock* mbb = mbb_ptr.get();
            VRegSet new_out;
            for (auto succ : mbb->successors) {
                new_out.insert(live_in[succ].begin(), live_in[succ].end());
            }
            VRegSet new_in = block_uses[mbb];
            VRegSet out_minus_def = new_out;
            for (unsigned d : block_defs[mbb]) out_minus_def.erase(d);
            new_in.insert(out_minus_def.begin(), out_minus_def.end());
            if (live_out[mbb] != new_out || live_in[mbb] != new_in) {
                changed = true;
                live_out[mbb] = new_out;
                live_in[mbb] = new_in;
            }
        }
    }

    for (const auto& mbb_ptr : MFunc->getBlocks()) {
        const MachineBasicBlock* mbb = mbb_ptr.get();
        VRegSet current_live = live_out[mbb];
        for (auto instr_it = mbb->getInstructions().rbegin(); instr_it != mbb->getInstructions().rend(); ++instr_it) {
            const MachineInstr* instr = instr_it->get();
            live_out_map[instr] = current_live;
            VRegSet use, def;
            // 使用新的、能看到物理寄存器的版本
            getInstrUseDef_Liveness(instr, use, def);
            for(auto d : def) current_live.erase(d);
            for(auto u : use) current_live.insert(u);
            live_in_map[instr] = current_live;
        }
    }
}

// [最终修正] 包含了所有正确逻辑和您已有调试代码的完整版本
void RISCv64RegAlloc::build() {
    initial.clear(); 
    RISCv64AsmPrinter printer_inside_build(MFunc);
    printer_inside_build.setStream(std::cerr);

    for (const auto& mbb_ptr : MFunc->getBlocks()) {
        for (const auto& instr_ptr : mbb_ptr->getInstructions()) {
            const MachineInstr* instr = instr_ptr.get();
            VRegSet use, def;
            getInstrUseDef_Liveness(instr, use, def);

            // 调试输出 use 和 def
            if (DEEPDEBUG) {
                std::cerr << "Instr:";
                printer_inside_build.printInstruction(instr_ptr.get(), true);
                // 修改 lambda 以捕获 this 指针，从而可以调用成员函数
                auto print_set = [this](const VRegSet& s, const std::string& name) {
                    std::cerr << "    " << name << ": { ";
                    for(unsigned v : s) std::cerr << regIdToString(v) << " "; // 使用新函数
                    std::cerr << "}\n";
                };
                print_set(def, "Def     ");
                print_set(use, "Use     ");
            }

            for (unsigned v : use) {
                if (!coloredNodes.count(v)) {
                    initial.insert(v);
                } else if (DEEPDEBUG) {
                    std::cerr << "Skipping %vreg" << v << " because it is in coloredNodes.\n";
                }
            }
            for (unsigned v : def) {
                if (!coloredNodes.count(v)) {
                    initial.insert(v);
                } else if (DEEPDEBUG) {
                    std::cerr << "Skipping %vreg" << v << " because it is in coloredNodes.\n";
                }
            }
        }
    }

    if (DEEPDEBUG) {
        std::cerr << "Initial set after build: { ";
        for (unsigned v : initial) std::cerr << "%vreg" << v << " ";
        std::cerr << "}\n";
    }
    
    for(unsigned vreg : initial) {
        adjList[vreg] = {};
        degree[vreg] = 0;
        moveList[vreg] = {};
        // alias[vreg] = vreg;
    }

    for (const auto& mbb_ptr : MFunc->getBlocks()) {
        if (DEEPDEBUG) std::cerr << "\n--- Building Graph for Basic Block: " << mbb_ptr->getName() << " ---\n";
        for (const auto& instr_ptr : mbb_ptr->getInstructions()) {
            const MachineInstr* instr = instr_ptr.get();
            VRegSet use, def;
            getInstrUseDef_Liveness(instr, use, def);
            const VRegSet& live_out = live_out_map.at(instr);

            if (DEEPDEBUG) {
                // 使用临时的 AsmPrinter 打印当前指令，便于观察
                RISCv64AsmPrinter temp_printer(MFunc);
                temp_printer.setStream(std::cerr);
                std::cerr << "Instr: ";
                temp_printer.printInstruction(const_cast<MachineInstr*>(instr), true);

                auto print_set = [this](const VRegSet& s, const std::string& name) {
                    std::cerr << "    " << name << ": { ";
                    for(unsigned v : s) std::cerr << regIdToString(v) << " "; // 使用新函数
                    std::cerr << "}\n";
                };
                print_set(def, "Def     ");
                print_set(use, "Use     ");
                print_set(live_out, "Live_Out");
                std::cerr << "  ----------------\n";
            }
            
            bool is_move = instr->getOpcode() == RVOpcodes::MV;

            if (is_move) {
                worklistMoves.insert(instr);
                VRegSet move_vregs;
                 for(const auto& op : instr->getOperands()) {
                    if (op->getKind() == MachineOperand::KIND_REG) {
                         auto reg_op = static_cast<const RegOperand*>(op.get());
                         if(reg_op->isVirtual()) move_vregs.insert(reg_op->getVRegNum());
                    }
                }
                for (unsigned vreg : move_vregs) {
                    moveList[vreg].insert(instr);
                }
            }

            // --- 规则 1: Def 与 Live_Out 变量干扰 ---
            VRegSet live = live_out;
            if (is_move) {
                for (unsigned u : use) {
                    live.erase(u);
                }
            }
            for (unsigned d : def) {
                for (unsigned l : live) {
                    addEdge(d, l);
                }
            }
            
            // --- [新增的关键逻辑] 规则 2: 对于非move指令，强制def和use互相干扰 ---
            // 这可以防止指令内部的源寄存器被目标寄存器错误地覆盖
            if (!is_move) {
                for (unsigned d : def) {
                    for (unsigned u : use) {
                        addEdge(d, u);
                    }
                }
            }

            // --- 规则 3: Live_Out 集合内部形成完全图 ---
            for (unsigned l1 : live_out) {
                for (unsigned l2 : live_out) {
                    addEdge(l1, l2);
                }
            }
        }
    }
}

// 将节点放入初始工作列表
void RISCv64RegAlloc::makeWorklist() {
    for (unsigned n : initial) {
        int K = isFPVReg(n) ? K_fp : K_int;
        if (degree.count(n) == 0) {
            std::cerr << "Error: degree not initialized for %vreg" << n << "\n";
            continue;
        }
        if (DEEPDEBUG) {
            std::cerr << "Assigning %vreg" << n << " (degree=" << degree.at(n)
                      << ", moveRelated=" << moveRelated(n) << ")\n";
        }
        if (degree.at(n) >= K) {
            spillWorklist.insert(n);
        } else if (moveRelated(n)) {
            freezeWorklist.insert(n);
        } else {
            simplifyWorklist.insert(n);
        }
    }
    if (DEEPDEBUG) std::cerr << "--------------------------------\n";
    initial.clear();
}

// 简化阶段
void RISCv64RegAlloc::simplify() {
    unsigned n = *simplifyWorklist.begin();
    simplifyWorklist.erase(simplifyWorklist.begin());
    if (DEEPDEBUG) std::cerr << "[Simplify] Popped %vreg" << n << ", pushing to stack.\n";
    selectStack.push_back(n);
    for (unsigned m : adjacent(n)) {
        decrementDegree(m);
    }
}

// 合并阶段
void RISCv64RegAlloc::coalesce() {
    const MachineInstr* move = *worklistMoves.begin();
    worklistMoves.erase(worklistMoves.begin());
    VRegSet use, def;
    getInstrUseDef_Liveness(move, use, def);
    unsigned x = getAlias(*def.begin());
    unsigned y = getAlias(*use.begin());
    unsigned u, v;
    if (precolored.count(y)) { u = y; v = x; } else { u = x; v = y; }

    if (DEEPDEBUG) std::cerr << "[Coalesce] Processing move between " << regIdToString(x) 
                           << " and " << regIdToString(y) << " (aliases " << regIdToString(u) 
                           << ", " << regIdToString(v) << ").\n";

    if (u == v) {
        if (DEEPDEBUG) std::cerr << "  -> Trivial coalesce (u == v).\n";
        coalescedMoves.insert(move);
        addWorklist(u);
        return; // 处理完毕，提前返回
    }

    if (isFPVReg(u) != isFPVReg(v)) {
        if (DEEPDEBUG) std::cerr << "  -> Constrained (type mismatch: " << regIdToString(u) << " is " 
                                << (isFPVReg(u) ? "float" : "int") << ", " << regIdToString(v) << " is "
                                << (isFPVReg(v) ? "float" : "int") << ").\n";
        constrainedMoves.insert(move);
        addWorklist(u);
        addWorklist(v);
        return; // 立即返回，不再进行后续检查
    }

    // --- 新的、拆分后的启发式检查逻辑 ---
    
    bool pre_interfere = adjList.at(v).count(u);
    
    if (pre_interfere) {
        if (DEEPDEBUG) std::cerr << "  -> Constrained (nodes already interfere).\n";
        constrainedMoves.insert(move);
        addWorklist(u);
        addWorklist(v);
        return;
    }

    bool is_u_precolored = precolored.count(u);
    bool can_coalesce = false;
    
    if (is_u_precolored) {
        // --- 场景1：u是物理寄存器，使用 George 启发式 ---
        if (DEEPDEBUG) std::cerr << "  -> Trying George Heuristic (u is precolored)...\n";
        
        // ==================== [展开的 std::all_of 逻辑] ====================
        
        // 步骤 1: 独立调用 adjacent(v) 获取邻居集合
        VRegSet neighbors_of_v = adjacent(v);
        if (DEEPDEBUG) {
            std::cerr << "      - Neighbors of " << regIdToString(v) << " to check are (" << neighbors_of_v.size() << "): { ";
            for (unsigned id : neighbors_of_v) std::cerr << regIdToString(id) << " ";
            std::cerr << "}\n";
        }

        // 步骤 2: 使用显式的 for 循环来代替 std::all_of
        bool george_ok = true; // 默认假设成功，任何一个邻居失败都会将此设为 false
        for (unsigned t : neighbors_of_v) {
            if (DEEPDEBUG) {
                std::cerr << "      - Checking neighbor " << regIdToString(t) << ":\n";
            }

            // 步骤 3: 独立调用启发式函数
            bool heuristic_result = georgeHeuristic(t, u);
            
            if (DEEPDEBUG) {
                std::cerr << "          - georgeHeuristic(" << regIdToString(t) << ", " << regIdToString(u) << ") -> " << (heuristic_result ? "OK" : "FAIL") << "\n";
            }

            if (!heuristic_result) {
                george_ok = false; // 只要有一个邻居不满足条件，整个检查就失败
                break;             // 并且可以立即停止检查其他邻居
            }
        }
        
        if (DEEPDEBUG) {
             std::cerr << "  -> George Heuristic final result: " << (george_ok ? "OK" : "FAIL") << "\n";
        }
        // =================================================================

        if (george_ok) {
            can_coalesce = true;
        }

    } else {
        // --- 场景2：u和v都是虚拟寄存器，使用 Briggs 启发式 ---
        if (DEEPDEBUG) std::cerr << "  -> Trying Briggs Heuristic (u and v are virtual)...\n";
        
        bool briggs_ok = briggsHeuristic(u, v);
        if (DEEPDEBUG) std::cerr << "      - briggsHeuristic(" << regIdToString(u) << ", " << regIdToString(v) << ") -> " << (briggs_ok ? "OK" : "FAIL") << "\n";

        if (briggs_ok) {
            can_coalesce = true;
        }
    }

    // --- 根据启发式结果进行最终决策 ---
    
    if (can_coalesce) {
        if (DEEPDEBUG) std::cerr << "  -> Heuristic OK. Combining " << regIdToString(v) << " into " << regIdToString(u) << ".\n";
        coalescedMoves.insert(move);
        combine(u, v);
        addWorklist(u);
    } else {
        if (DEEPDEBUG) std::cerr << "  -> Heuristic failed. Adding to active moves.\n";
        activeMoves.insert(move);
    }
}

// 冻结阶段
void RISCv64RegAlloc::freeze() {
    unsigned u = *freezeWorklist.begin();
    freezeWorklist.erase(freezeWorklist.begin());
    if (DEEPDEBUG) std::cerr << "[Freeze] Freezing %vreg" << u << " and moving to simplify list.\n";
    simplifyWorklist.insert(u);
    freezeMoves(u);
}

// 选择溢出节点
void RISCv64RegAlloc::selectSpill() {
    auto it = std::max_element(spillWorklist.begin(), spillWorklist.end(), 
        [&](unsigned a, unsigned b){ return degree.at(a) < degree.at(b); });
    unsigned m = *it;
    spillWorklist.erase(it);
    if (DEEPDEBUG) std::cerr << "[Spill] Selecting %vreg" << m << " to spill.\n";
    simplifyWorklist.insert(m);
    freezeMoves(m);
}

void RISCv64RegAlloc::assignColors() {
    if (DEEPDEBUG) std::cerr << "[AssignColors] Starting...\n";
    while (!selectStack.empty()) {
        unsigned n = selectStack.back();
        selectStack.pop_back();
        bool is_fp = isFPVReg(n);
        const auto& available_regs = is_fp ? allocable_fp_regs : allocable_int_regs;
        std::set<PhysicalReg> ok_colors(available_regs.begin(), available_regs.end());

        // 遍历 n 的所有邻居 w
        for (unsigned w : adjList.at(n)) {
            unsigned w_alias = getAlias(w);

            // [关键修正] 邻居 w 可能是一个已着色的虚拟寄存器，
            // 或者它本身就是一个物理寄存器。两种情况都要处理！
            bool is_colored_vreg = coloredNodes.count(w_alias);
            bool is_physical_reg = precolored.count(w_alias);
            const unsigned offset = static_cast<unsigned>(PhysicalReg::PHYS_REG_START_ID);

            if (is_colored_vreg || is_physical_reg) {
                PhysicalReg neighbor_color;
                if (is_colored_vreg) {
                    // 如果是已着色的vreg，从 color_map 获取它的颜色
                    neighbor_color = color_map.at(w_alias);
                } else {
                    // 如果是物理寄存器，它的ID就是它的颜色
                    neighbor_color = static_cast<PhysicalReg>(w_alias - offset);
                }
                ok_colors.erase(neighbor_color);
            }
        }

        if (ok_colors.empty()) {
            if (DEEPDEBUG) std::cerr << "  -> No color for %vreg" << n << ". Spilling.\n";
            spilledNodes.insert(n);
        } else {
            PhysicalReg c = *ok_colors.begin();
            coloredNodes.insert(n);
            color_map[n] = c;
            if (DEEPDEBUG) std::cerr << "  -> Colored %vreg" << n << " with " << regToString(c) << " (ID: " << static_cast<int>(c) << ").\n";
        }
    }

    // 为合并的节点上色（这部分逻辑是正确的）
    for (unsigned n : coalescedNodes) {
        unsigned root_alias = getAlias(n);
        if (color_map.count(root_alias)) {
            color_map[n] = color_map.at(root_alias);
            if (DEEPDEBUG) std::cerr << "  -> Coalesced " << regIdToString(n) << " gets color of alias " << regIdToString(root_alias) << ".\n";
        } else {
             if (DEEPDEBUG) std::cerr << "  -> No color for alias of %vreg" << n << ". Spilling.\n";
            spilledNodes.insert(n);
        }
    }
}

// 重写程序，插入溢出代码
void RISCv64RegAlloc::rewriteProgram() {
    // 1. 为溢出的旧vreg在栈上分配空间 (这部分逻辑不变)
    StackFrameInfo& frame_info = MFunc->getFrameInfo();
    // locals_size, callee_saved_size等已经由之前的步骤或上一次运行计算好
    // 我们在此基础上累加溢出区大小
    int spill_base_offset = frame_info.locals_size + frame_info.callee_saved_size;

    for (unsigned vreg : spilledNodes) {
        if (frame_info.spill_offsets.count(vreg)) continue; // 如果已经分配过则跳过

        int size = 4; // 默认4字节
        if (isFPVReg(vreg)) {
            size = 4; // float
        } else if (vreg_type_map.count(vreg) && vreg_type_map.at(vreg)->isPointer()) {
            size = 8; // pointer
        }

        spill_base_offset += size;
        spill_base_offset = (spill_base_offset + 7) & ~7; // 8字节对齐
        frame_info.spill_offsets[vreg] = spill_base_offset;
    }
    frame_info.spill_size = spill_base_offset - (frame_info.locals_size + frame_info.callee_saved_size);

    // 2. 遍历所有指令，为溢出vreg的use/def插入load/store，并使用新的临时虚拟寄存器
    for (auto& mbb : MFunc->getBlocks()) {
        std::vector<std::unique_ptr<MachineInstr>> new_instructions;
        
        for (auto& instr_ptr : mbb->getInstructions()) {
            // 对每条指令，记录其使用的旧vreg到新临时vreg的映射
            std::map<unsigned, unsigned> vreg_remap; 

            // a. 为每个溢出的 use 操作数，在原指令前插入 load 指令
            VRegSet use, def;
            getInstrUseDef_Liveness(instr_ptr.get(), use, def); // 假设此函数能正确处理operands

            for(auto& op : instr_ptr->getOperands()) {
                if (op->getKind() == MachineOperand::KIND_REG) {
                    auto reg_op = static_cast<RegOperand*>(op.get());
                    if (reg_op->isVirtual() && spilledNodes.count(reg_op->getVRegNum()) && use.count(reg_op->getVRegNum())) {
                        unsigned old_vreg = reg_op->getVRegNum();
                        // 仅当还未为此vreg创建临时替身时才创建
                        if (vreg_remap.find(old_vreg) == vreg_remap.end()) {
                            Type* type = vreg_type_map.count(old_vreg) ? vreg_type_map.at(old_vreg) : Type::getIntType();
                            unsigned new_temp_vreg = ISel->getNewVReg(type);
                            vreg_remap[old_vreg] = new_temp_vreg;

                            RVOpcodes load_op;
                            if (isFPVReg(old_vreg)) load_op = RVOpcodes::FLW;
                            else if (type->isPointer()) load_op = RVOpcodes::LD;
                            else load_op = RVOpcodes::LW;

                            auto load = std::make_unique<MachineInstr>(load_op);
                            load->addOperand(std::make_unique<RegOperand>(new_temp_vreg));
                            load->addOperand(std::make_unique<MemOperand>(
                                std::make_unique<RegOperand>(PhysicalReg::S0),
                                std::make_unique<ImmOperand>(-frame_info.spill_offsets.at(old_vreg)) // 偏移是相对于S0的负值
                            ));
                            new_instructions.push_back(std::move(load));
                        }
                    }
                } else if (op->getKind() == MachineOperand::KIND_MEM) {
                    auto reg_op = static_cast<MemOperand*>(op.get())->getBase();
                    if (reg_op->isVirtual() && spilledNodes.count(reg_op->getVRegNum())) {
                         unsigned old_vreg = reg_op->getVRegNum();
                         if (vreg_remap.find(old_vreg) == vreg_remap.end()) {
                            Type* type = vreg_type_map.count(old_vreg) ? vreg_type_map.at(old_vreg) : Type::getIntType();
                            unsigned new_temp_vreg = ISel->getNewVReg(type);
                            vreg_remap[old_vreg] = new_temp_vreg;

                            RVOpcodes load_op = (type->isPointer()) ? RVOpcodes::LD : RVOpcodes::LW;
                            auto load = std::make_unique<MachineInstr>(load_op);
                            load->addOperand(std::make_unique<RegOperand>(new_temp_vreg));
                            load->addOperand(std::make_unique<MemOperand>(
                                std::make_unique<RegOperand>(PhysicalReg::S0),
                                std::make_unique<ImmOperand>(-frame_info.spill_offsets.at(old_vreg))
                            ));
                            new_instructions.push_back(std::move(load));
                         }
                    }
                }
            }
            
            // b. 替换原指令中的操作数，并将其加入新指令列表
            for(auto& op : instr_ptr->getOperands()) {
                if (op->getKind() == MachineOperand::KIND_REG) {
                    auto reg_op = static_cast<RegOperand*>(op.get());
                    if (reg_op->isVirtual() && vreg_remap.count(reg_op->getVRegNum())) {
                        reg_op->setVRegNum(vreg_remap.at(reg_op->getVRegNum()));
                    }
                } else if (op->getKind() == MachineOperand::KIND_MEM) {
                     auto reg_op = static_cast<MemOperand*>(op.get())->getBase();
                     if (reg_op->isVirtual() && vreg_remap.count(reg_op->getVRegNum())) {
                        reg_op->setVRegNum(vreg_remap.at(reg_op->getVRegNum()));
                    }
                }
            }
            new_instructions.push_back(std::move(instr_ptr));

            // c. 为每个溢出的 def 操作数，在原指令后插入 store 指令
            vreg_remap.clear(); // 清空映射，为def创建新的临时vreg
            for(auto& op : new_instructions.back()->getOperands()) {
                 if (op->getKind() == MachineOperand::KIND_REG) {
                    auto reg_op = static_cast<RegOperand*>(op.get());
                     if (reg_op->isVirtual() && spilledNodes.count(reg_op->getVRegNum()) && def.count(reg_op->getVRegNum())) {
                        unsigned old_vreg = reg_op->getVRegNum();
                        if (vreg_remap.find(old_vreg) == vreg_remap.end()) {
                            Type* type = vreg_type_map.count(old_vreg) ? vreg_type_map.at(old_vreg) : Type::getIntType();
                            unsigned new_temp_vreg = ISel->getNewVReg(type);
                            vreg_remap[old_vreg] = new_temp_vreg; // 记录从old到new的映射
                            
                            // 再次修改指令，将def的vreg也换成新的临时vreg
                            reg_op->setVRegNum(new_temp_vreg);

                            RVOpcodes store_op;
                            if (isFPVReg(old_vreg)) store_op = RVOpcodes::FSW;
                            else if (type->isPointer()) store_op = RVOpcodes::SD;
                            else store_op = RVOpcodes::SW;

                            auto store = std::make_unique<MachineInstr>(store_op);
                            store->addOperand(std::make_unique<RegOperand>(new_temp_vreg)); // 从新的临时vreg中存值
                            store->addOperand(std::make_unique<MemOperand>(
                                std::make_unique<RegOperand>(PhysicalReg::S0),
                                std::make_unique<ImmOperand>(-frame_info.spill_offsets.at(old_vreg))
                            ));
                            // 在原指令之后插入store
                            new_instructions.push_back(std::move(store));
                        }
                    }
                 }
            }
        }
        mbb->getInstructions() = std::move(new_instructions);
    }
    
    // 清空溢出节点集合，为下一次迭代分配做准备
    spilledNodes.clear();
}

/**
 * @brief [最终修正] 获取一条指令完整的【虚拟】使用/定义寄存器集合
 * 这个函数将服务于图的构建（收集initial节点等）。
 */
void RISCv64RegAlloc::getInstrUseDef(const MachineInstr* instr, VRegSet& use, VRegSet& def) {
    auto opcode = instr->getOpcode();
    const auto& operands = instr->getOperands();

    static const std::map<RVOpcodes, std::pair<std::vector<int>, std::vector<int>>> op_info = {
        {RVOpcodes::ADD, {{0}, {1, 2}}}, {RVOpcodes::SUB, {{0}, {1, 2}}}, {RVOpcodes::MUL, {{0}, {1, 2}}},
        {RVOpcodes::DIV, {{0}, {1, 2}}}, {RVOpcodes::REM, {{0}, {1, 2}}}, {RVOpcodes::ADDW, {{0}, {1, 2}}},
        {RVOpcodes::SUBW, {{0}, {1, 2}}}, {RVOpcodes::MULW, {{0}, {1, 2}}}, {RVOpcodes::DIVW, {{0}, {1, 2}}},
        {RVOpcodes::REMW, {{0}, {1, 2}}}, {RVOpcodes::SLT, {{0}, {1, 2}}}, {RVOpcodes::SLTU, {{0}, {1, 2}}},
        {RVOpcodes::ADDI, {{0}, {1}}}, {RVOpcodes::ADDIW, {{0}, {1}}}, {RVOpcodes::XORI, {{0}, {1}}},
        {RVOpcodes::SLTI, {{0}, {1}}}, {RVOpcodes::SLTIU, {{0}, {1}}}, {RVOpcodes::LB, {{0}, {}}},
        {RVOpcodes::LH, {{0}, {}}}, {RVOpcodes::LW, {{0}, {}}}, {RVOpcodes::LD, {{0}, {}}},
        {RVOpcodes::LBU, {{0}, {}}}, {RVOpcodes::LHU, {{0}, {}}}, {RVOpcodes::LWU, {{0}, {}}},
        {RVOpcodes::FLW, {{0}, {}}}, {RVOpcodes::FLD, {{0}, {}}}, {RVOpcodes::SB, {{}, {0, 1}}},
        {RVOpcodes::SH, {{}, {0, 1}}}, {RVOpcodes::SW, {{}, {0, 1}}}, {RVOpcodes::SD, {{}, {0, 1}}},
        {RVOpcodes::FSW, {{}, {0, 1}}}, {RVOpcodes::FSD, {{}, {0, 1}}}, {RVOpcodes::BEQ, {{}, {0, 1}}},
        {RVOpcodes::BNE, {{}, {0, 1}}}, {RVOpcodes::BLT, {{}, {0, 1}}}, {RVOpcodes::BGE, {{}, {0, 1}}},
        {RVOpcodes::JALR, {{0}, {1}}}, {RVOpcodes::LI, {{0}, {}}}, {RVOpcodes::LA, {{0}, {}}},
        {RVOpcodes::MV, {{0}, {1}}}, {RVOpcodes::SEQZ, {{0}, {1}}}, {RVOpcodes::SNEZ, {{0}, {1}}},
        {RVOpcodes::RET, {{}, {}}}, {RVOpcodes::FADD_S, {{0}, {1, 2}}}, {RVOpcodes::FSUB_S, {{0}, {1, 2}}},
        {RVOpcodes::FMUL_S, {{0}, {1, 2}}}, {RVOpcodes::FDIV_S, {{0}, {1, 2}}}, {RVOpcodes::FEQ_S, {{0}, {1, 2}}},
        {RVOpcodes::FLT_S, {{0}, {1, 2}}}, {RVOpcodes::FLE_S, {{0}, {1, 2}}}, {RVOpcodes::FCVT_S_W, {{0}, {1}}},
        {RVOpcodes::FCVT_W_S, {{0}, {1}}}, {RVOpcodes::FMV_S, {{0}, {1}}}, {RVOpcodes::FMV_W_X, {{0}, {1}}},
        {RVOpcodes::FMV_X_W, {{0}, {1}}}, {RVOpcodes::FNEG_S, {{0}, {1}}}
    };
    
    auto get_vreg_id_if_virtual = [&](const MachineOperand* op, VRegSet& s) {
        if (op->getKind() == MachineOperand::KIND_REG) {
            auto reg_op = static_cast<const RegOperand*>(op);
            if (reg_op->isVirtual()) s.insert(reg_op->getVRegNum());
        } else if (op->getKind() == MachineOperand::KIND_MEM) {
            auto mem_op = static_cast<const MemOperand*>(op);
            auto reg_op = mem_op->getBase();
            if (reg_op->isVirtual()) s.insert(reg_op->getVRegNum());
        }
    };

    if (op_info.count(opcode)) {
        const auto& info = op_info.at(opcode);
        for (int idx : info.first) if (idx < operands.size()) get_vreg_id_if_virtual(operands[idx].get(), def);
        for (int idx : info.second) if (idx < operands.size()) get_vreg_id_if_virtual(operands[idx].get(), use);
        for (const auto& op : operands) if (op->getKind() == MachineOperand::KIND_MEM) get_vreg_id_if_virtual(op.get(), use);
    } else if (opcode == RVOpcodes::CALL) {
        if (!operands.empty() && operands[0]->getKind() == MachineOperand::KIND_REG) get_vreg_id_if_virtual(operands[0].get(), def);
        for (size_t i = 1; i < operands.size(); ++i) if (operands[i]->getKind() == MachineOperand::KIND_REG) get_vreg_id_if_virtual(operands[i].get(), use);
    }
}

/**
 * @brief [新增] 获取一条指令完整的、包含物理寄存器的Use/Def集合
 * 这个新函数将专门服务于活跃性分析。
 */
void RISCv64RegAlloc::getInstrUseDef_Liveness(const MachineInstr* instr, VRegSet& use, VRegSet& def) {
    auto opcode = instr->getOpcode();
    const auto& operands = instr->getOperands();

    static const std::map<RVOpcodes, std::pair<std::vector<int>, std::vector<int>>> op_info = {
        {RVOpcodes::ADD, {{0}, {1, 2}}}, {RVOpcodes::SUB, {{0}, {1, 2}}}, {RVOpcodes::MUL, {{0}, {1, 2}}},
        {RVOpcodes::DIV, {{0}, {1, 2}}}, {RVOpcodes::REM, {{0}, {1, 2}}}, {RVOpcodes::ADDW, {{0}, {1, 2}}},
        {RVOpcodes::SUBW, {{0}, {1, 2}}}, {RVOpcodes::MULW, {{0}, {1, 2}}}, {RVOpcodes::DIVW, {{0}, {1, 2}}},
        {RVOpcodes::REMW, {{0}, {1, 2}}}, {RVOpcodes::SLT, {{0}, {1, 2}}}, {RVOpcodes::SLTU, {{0}, {1, 2}}},
        {RVOpcodes::ADDI, {{0}, {1}}}, {RVOpcodes::ADDIW, {{0}, {1}}}, {RVOpcodes::XORI, {{0}, {1}}},
        {RVOpcodes::SLTI, {{0}, {1}}}, {RVOpcodes::SLTIU, {{0}, {1}}}, {RVOpcodes::LB, {{0}, {}}},
        {RVOpcodes::LH, {{0}, {}}}, {RVOpcodes::LW, {{0}, {}}}, {RVOpcodes::LD, {{0}, {}}},
        {RVOpcodes::LBU, {{0}, {}}}, {RVOpcodes::LHU, {{0}, {}}}, {RVOpcodes::LWU, {{0}, {}}},
        {RVOpcodes::FLW, {{0}, {}}}, {RVOpcodes::FLD, {{0}, {}}}, {RVOpcodes::SB, {{}, {0, 1}}},
        {RVOpcodes::SH, {{}, {0, 1}}}, {RVOpcodes::SW, {{}, {0, 1}}}, {RVOpcodes::SD, {{}, {0, 1}}},
        {RVOpcodes::FSW, {{}, {0, 1}}}, {RVOpcodes::FSD, {{}, {0, 1}}}, {RVOpcodes::BEQ, {{}, {0, 1}}},
        {RVOpcodes::BNE, {{}, {0, 1}}}, {RVOpcodes::BLT, {{}, {0, 1}}}, {RVOpcodes::BGE, {{}, {0, 1}}},
        {RVOpcodes::JALR, {{0}, {1}}}, {RVOpcodes::LI, {{0}, {}}}, {RVOpcodes::LA, {{0}, {}}},
        {RVOpcodes::MV, {{0}, {1}}}, {RVOpcodes::SEQZ, {{0}, {1}}}, {RVOpcodes::SNEZ, {{0}, {1}}},
        {RVOpcodes::RET, {{}, {static_cast<int>(PhysicalReg::A0), static_cast<int>(PhysicalReg::F10)}}},
        {RVOpcodes::FADD_S, {{0}, {1, 2}}}, {RVOpcodes::FSUB_S, {{0}, {1, 2}}}, {RVOpcodes::FMUL_S, {{0}, {1, 2}}},
        {RVOpcodes::FDIV_S, {{0}, {1, 2}}}, {RVOpcodes::FEQ_S, {{0}, {1, 2}}}, {RVOpcodes::FLT_S, {{0}, {1, 2}}},
        {RVOpcodes::FLE_S, {{0}, {1, 2}}}, {RVOpcodes::FCVT_S_W, {{0}, {1}}}, {RVOpcodes::FCVT_W_S, {{0}, {1}}},
        {RVOpcodes::FMV_S, {{0}, {1}}}, {RVOpcodes::FMV_W_X, {{0}, {1}}}, {RVOpcodes::FMV_X_W, {{0}, {1}}},
        {RVOpcodes::FNEG_S, {{0}, {1}}}
    };
    
    const unsigned offset = static_cast<unsigned>(PhysicalReg::PHYS_REG_START_ID);
    auto get_any_reg_id = [&](const MachineOperand* op) -> unsigned {
        if (op->getKind() == MachineOperand::KIND_REG) {
            auto reg_op = static_cast<const RegOperand*>(op);
            return reg_op->isVirtual() ? reg_op->getVRegNum() : (offset + static_cast<unsigned>(reg_op->getPReg()));
        } else if (op->getKind() == MachineOperand::KIND_MEM) {
            auto mem_op = static_cast<const MemOperand*>(op);
            auto reg_op = mem_op->getBase();
            return reg_op->isVirtual() ? reg_op->getVRegNum() : (offset + static_cast<unsigned>(reg_op->getPReg()));
        }
        return (unsigned)-1;
    };

    if (op_info.count(opcode)) {
        const auto& info = op_info.at(opcode);
        for (int idx : info.first) if (idx < operands.size()) def.insert(get_any_reg_id(operands[idx].get()));
        for (int idx : info.second) if (idx < operands.size()) use.insert(get_any_reg_id(operands[idx].get()));
        for (const auto& op : operands) if (op->getKind() == MachineOperand::KIND_MEM) use.insert(get_any_reg_id(op.get()));
    } else if (opcode == RVOpcodes::CALL) {
        if (!operands.empty() && operands[0]->getKind() == MachineOperand::KIND_REG) def.insert(get_any_reg_id(operands[0].get()));
        for (size_t i = 1; i < operands.size(); ++i) if (operands[i]->getKind() == MachineOperand::KIND_REG) use.insert(get_any_reg_id(operands[i].get()));
        for (auto preg : getCallerSavedIntRegs()) def.insert(static_cast<unsigned>(preg));
        for (auto preg : getCallerSavedFpRegs()) def.insert(static_cast<unsigned>(preg));
        def.insert(static_cast<unsigned>(PhysicalReg::RA));
    }
}

void RISCv64RegAlloc::addEdge(unsigned u, unsigned v) {
    if (u == v) return;

    // 关键修正：只为虚拟寄存器（非预着色）更新邻接表和度数
    // 如果 u 是虚拟寄存器
    if (!precolored.count(u)) {
        // 并且 u 和 v 之间还没有边
        if (adjList.at(u).find(v) == adjList.at(u).end()) {
            adjList.at(u).insert(v);
            degree.at(u)++;
        }
    }

    // 对称地，如果 v 是虚拟寄存器
    if (!precolored.count(v)) {
        // 并且 v 和 u 之间还没有边
        if (adjList.at(v).find(u) == adjList.at(v).end()) {
            adjList.at(v).insert(u);
            degree.at(v)++;
        }
    }
}

RISCv64RegAlloc::VRegSet RISCv64RegAlloc::adjacent(unsigned n) {
    // 仅在 DEEPDEBUG 模式下启用详细日志
    if (DEEPDEBUG) {
        // 使用 regIdToString 打印节点 n，无论是物理还是虚拟
        std::cerr << "\n[adjacent] >>>>> Executing for node " << regIdToString(n) << " <<<<<\n";
    }

    // 1. 如果节点 n 是物理寄存器，它没有邻接表，直接返回空集
    if (precolored.count(n)) {
        if (DEEPDEBUG) {
            std::cerr << "[adjacent] Node " << regIdToString(n) << " is precolored. Returning {}.\n";
        }
        return {};
    }

    // 安全检查：确保 n 在 adjList 中存在，防止 map::at 崩溃
    if (adjList.count(n) == 0) {
        if (DEEPDEBUG) {
            std::cerr << "[adjacent] WARNING: Node " << regIdToString(n) << " not found in adjList. Returning {}.\n";
        }
        return {};
    }

    // 2. 获取 n 在冲突图中的所有邻居
    VRegSet result = adjList.at(n);
    
    if (DEEPDEBUG) {
        // 定义一个局部的 lambda 方便打印集合
        auto print_set = [this](const VRegSet& s, const std::string& name) {
            std::cerr << "[adjacent] " << name << " (" << s.size() << "): { ";
            for (unsigned id : s) std::cerr << regIdToString(id) << " ";
            std::cerr << "}\n";
        };
        print_set(result, "Initial full neighbors");
    }

    // 3. 过滤掉那些已经在 selectStack 或 coalescedNodes 中的邻居
    //    这些节点被认为是“已移除”的，不参与当前的启发式判断

    // 3a. 从 selectStack 中移除
    VRegSet removed_from_stack; // 仅用于调试打印
    for (auto it = selectStack.rbegin(); it != selectStack.rend(); ++it) {
        if (result.count(*it)) {
            if (DEEPDEBUG) removed_from_stack.insert(*it);
            result.erase(*it);
        }
    }
    if (DEEPDEBUG && !removed_from_stack.empty()) {
        std::cerr << "[adjacent]   - Removed from selectStack: { ";
        for(unsigned id : removed_from_stack) std::cerr << regIdToString(id) << " ";
        std::cerr << "}\n";
    }

    // 3b. 从 coalescedNodes 中移除
    VRegSet removed_from_coalesced; // 仅用于调试打印
    for (unsigned cn : coalescedNodes) {
        if (result.count(cn)) {
            if (DEEPDEBUG) removed_from_coalesced.insert(cn);
            result.erase(cn);
        }
    }
    if (DEEPDEBUG && !removed_from_coalesced.empty()) {
        std::cerr << "[adjacent]   - Removed from coalescedNodes: { ";
        for(unsigned id : removed_from_coalesced) std::cerr << regIdToString(id) << " ";
        std::cerr << "}\n";
    }

    // 4. 返回最终的、过滤后的“有效”邻居集合
    if (DEEPDEBUG) {
        std::cerr << "[adjacent] >>>>> Returning final adjacent set (" << result.size() << "): { ";
        for (unsigned id : result) std::cerr << regIdToString(id) << " ";
        std::cerr << "}\n\n";
    }

    return result;
}

RISCv64RegAlloc::VRegMoveSet RISCv64RegAlloc::nodeMoves(unsigned n) {
    if (precolored.count(n) || !moveList.count(n)) {
        return {};
    }
    
    VRegMoveSet result;
    const VRegMoveSet& moves = moveList.at(n);
    for (const auto& move : moves) {
        if (activeMoves.count(move) || worklistMoves.count(move)) {
            result.insert(move);
        }
    }
    return result;
}

bool RISCv64RegAlloc::moveRelated(unsigned n) {
    return !nodeMoves(n).empty();
}

void RISCv64RegAlloc::decrementDegree(unsigned m) {
    if (precolored.count(m)) {
        return;
    }

    int K = isFPVReg(m) ? K_fp : K_int;
    int d = degree.at(m);
    degree.at(m)--;
    if (d == K) {
        VRegSet nodes_to_enable = adjacent(m);
        nodes_to_enable.insert(m);
        enableMoves(nodes_to_enable);
        spillWorklist.erase(m);
        if (moveRelated(m)) {
            freezeWorklist.insert(m);
        } else {
            simplifyWorklist.insert(m);
        }
    }
}

void RISCv64RegAlloc::enableMoves(const VRegSet& nodes) {
    for (unsigned n : nodes) {
        VRegMoveSet moves = nodeMoves(n);
        for (const auto& move : moves) {
            if (activeMoves.count(move)) {
                activeMoves.erase(move);
                worklistMoves.insert(move);
            }
        }
    }
}

unsigned RISCv64RegAlloc::getAlias(unsigned n) {
    if (precolored.count(n)) {
        return n;
    }
    if (alias.count(n)) {
        // 路径压缩
        alias.at(n) = getAlias(alias.at(n));
        return alias.at(n);
    }
    return n;
}

void RISCv64RegAlloc::addWorklist(unsigned u) {
    if (precolored.count(u)) return;

    int K = isFPVReg(u) ? K_fp : K_int;
    if (!moveRelated(u) && degree.at(u) < K) {
        freezeWorklist.erase(u);
        simplifyWorklist.insert(u);
    }
}

// Briggs启发式
bool RISCv64RegAlloc::briggsHeuristic(unsigned u, unsigned v) {
    if (DEEPDEBUG) {
        std::cerr << "\n[Briggs] >>>>> Checking coalesce between " << regIdToString(u) << " and " << regIdToString(v) << " <<<<<\n";
    }

    // 步骤 1: 分别获取 u 和 v 的邻居
    VRegSet u_adj = adjacent(u);
    VRegSet v_adj = adjacent(v);

    // 步骤 2: 合并两个邻居集合
    VRegSet all_adj = u_adj;
    all_adj.insert(v_adj.begin(), v_adj.end());

    if (DEEPDEBUG) {
        auto print_set = [this](const VRegSet& s, const std::string& name) {
            std::cerr << "[Briggs] " << name << " (" << s.size() << "): { ";
            for (unsigned id : s) std::cerr << regIdToString(id) << " ";
            std::cerr << "}\n";
        };
        print_set(u_adj, "Neighbors of u");
        print_set(v_adj, "Neighbors of v");
        print_set(all_adj, "Combined neighbors");
    }

    // 步骤 3: 遍历合并后的邻居集合，计算度数 >= K 的节点数量
    int k = 0;
    if (DEEPDEBUG) std::cerr << "[Briggs] Checking significance of combined neighbors:\n";
    for (unsigned n : all_adj) {
        // 关键修正：只考虑那些在工作集中的邻居节点 n
        if (degree.count(n) > 0) {
            int K = isFPVReg(n) ? K_fp : K_int;
            if (degree.at(n) >= K) {
                k++;
                if (DEEPDEBUG) {
                    std::cerr << "[Briggs]   - Node " << regIdToString(n) << " is significant (degree " << degree.at(n) << " >= " << K << "). Count k is now " << k << ".\n";
                }
            }
        }
    }

    // 步骤 4: 比较 "重要" 邻居的数量是否小于 K
    int K_u = isFPVReg(u) ? K_fp : K_int;
    bool result = (k < K_u);

    if (DEEPDEBUG) {
        std::cerr << "[Briggs] Final count of significant neighbors (k) = " << k << ".\n";
        std::cerr << "[Briggs] K value for node " << regIdToString(u) << " is " << K_u << ".\n";
        std::cerr << "[Briggs] >>>>> Result (k < K): " << (result ? "OK (can coalesce)" : "FAIL (cannot coalesce)") << "\n\n";
    }

    return result;
}

// George启发式
bool RISCv64RegAlloc::georgeHeuristic(unsigned t, unsigned u) {
    // 关键修正：如果 t 不是一个待分配的虚拟寄存器（即它是物理寄存器），
    // 那么它已经被预着色，总是满足 George 启发式条件。
    // 我们通过检查 degree.count(t) 来判断 t 是否在我们的虚拟寄存器工作集中。
    if (degree.count(t) == 0) {
        return true;
    }
    
    int K = isFPVReg(t) ? K_fp : K_int;
    // adjList.at(t) 现在是安全的，因为 degree.count(t) > 0 保证了 adjList.count(t) > 0
    return degree.at(t) < K || precolored.count(u) || adjList.at(t).count(u);
}

void RISCv64RegAlloc::combine(unsigned u, unsigned v) {
    freezeWorklist.erase(v);
    spillWorklist.erase(v);
    coalescedNodes.insert(v);
    alias[v] = u;
    if (moveList.count(u) && moveList.count(v)) {
        moveList.at(u).insert(moveList.at(v).begin(), moveList.at(v).end());
    } else if (moveList.count(v)) {
        moveList[u] = moveList.at(v);
    }
    enableMoves({v});
    for (unsigned t : adjacent(v)) {
        addEdge(t, u);
        decrementDegree(t);
    }
    int K = isFPVReg(u) ? K_fp : K_int;
    if (degree.at(u) >= K && freezeWorklist.count(u)) {
        freezeWorklist.erase(u);
        spillWorklist.insert(u);
    }
}

void RISCv64RegAlloc::freezeMoves(unsigned u) {
    if (precolored.count(u)) return;

    VRegMoveSet moves = nodeMoves(u);
    for (const auto& move : moves) {
        VRegSet use, def;
        getInstrUseDef_Liveness(move, use, def);
        unsigned x = *def.begin();
        unsigned y = *use.begin();
        unsigned v_alias;

        if (getAlias(y) == getAlias(u)) {
            v_alias = getAlias(x);
        } else {
            v_alias = getAlias(y);
        }

        activeMoves.erase(move);
        frozenMoves.insert(move);

        if (!precolored.count(v_alias) && nodeMoves(v_alias).empty() && degree.at(v_alias) < (isFPVReg(v_alias) ? K_fp : K_int)) {
            freezeWorklist.erase(v_alias);
            simplifyWorklist.insert(v_alias);
        }
    }
}

// 检查vreg是否为浮点类型
bool RISCv64RegAlloc::isFPVReg(unsigned vreg) const {
    // 1. 检查是否为虚拟寄存器
    if (vreg_type_map.count(vreg)) {
        return vreg_type_map.at(vreg)->isFloat();
    }
    
    // 2. 检查是否为物理寄存器 (ID >= 100000)
    const unsigned offset = static_cast<unsigned>(PhysicalReg::PHYS_REG_START_ID);
    if (vreg >= offset) {
        // 先减去偏移，还原成原始的、小的枚举值
        unsigned raw_preg_id = vreg - offset;
        
        // 再用原始枚举值判断是否在浮点寄存器范围内
        if (raw_preg_id >= static_cast<unsigned>(PhysicalReg::F0) && raw_preg_id <= static_cast<unsigned>(PhysicalReg::F31)) {
            return true;
        }
    }

    // 3. 其他所有情况（如未知的vreg，或整数物理寄存器）都默认为整数
    return false;
}

// 收集被使用的被调用者保存寄存器
void RISCv64RegAlloc::collectUsedCalleeSavedRegs() {
    StackFrameInfo& frame_info = MFunc->getFrameInfo();
    frame_info.used_callee_saved_regs.clear();

    const auto& callee_saved_int = getCalleeSavedIntRegs();
    const auto& callee_saved_fp = getCalleeSavedFpRegs();
    std::set<PhysicalReg> callee_saved_set(callee_saved_int.begin(), callee_saved_int.end());
    callee_saved_set.insert(callee_saved_fp.begin(), callee_saved_fp.end());
    // s0总是被使用作为帧指针
    callee_saved_set.insert(PhysicalReg::S0);


    for(const auto& pair : color_map) {
        PhysicalReg preg = pair.second;
        if(callee_saved_set.count(preg)) {
            frame_info.used_callee_saved_regs.insert(preg);
        }
    }
}

/**
 * @brief 将最终的寄存器分配结果应用到所有机器指令上。
 * 遍历所有操作数，将虚拟寄存器替换为分配到的物理寄存器。
 */
void RISCv64RegAlloc::applyColoring() {
    for (auto& mbb : MFunc->getBlocks()) {
        for (auto& instr_ptr : mbb->getInstructions()) {
            for (auto& op_ptr : instr_ptr->getOperands()) {
                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)) {
                            // 使用 setPReg 将虚拟寄存器转换为物理寄存器
                            reg_op->setPReg(color_map.at(vreg));
                        } else {
                            // 如果一个vreg在成功分配后仍然没有颜色，这是一个错误
                            std::cerr << "FATAL: Virtual register %vreg" << vreg << " has no color after allocation!\n";
                            assert(false && "Virtual register has no color after allocation!");
                            reg_op->setPReg(PhysicalReg::T6);
                        }
                    }
                } else if (op_ptr->getKind() == MachineOperand::KIND_MEM) {
                    auto mem_op = static_cast<MemOperand*>(op_ptr.get());
                    auto reg_op = mem_op->getBase();
                    if (reg_op->isVirtual()) {
                        unsigned vreg = reg_op->getVRegNum();
                        if (color_map.count(vreg)) {
                            reg_op->setPReg(color_map.at(vreg));
                        } else {
                            assert(false && "Virtual register in memory operand has no color!");
                            reg_op->setPReg(PhysicalReg::T6);
                        }
                    }
                }
            }
        }
    }
}

void RISCv64RegAlloc::dumpState(const std::string& stage) {
    if (!DEEPDEBUG) return;
    std::cerr << "\n=============== STATE DUMP (" << stage << ") ===============\n";
    auto print_vreg_set = [&](const VRegSet& s, const std::string& name){
        std::cerr << name << " (" << s.size() << "): { ";
        for(unsigned v : s) std::cerr << "%vreg" << v << " ";
        std::cerr << "}\n";
    };
    print_vreg_set(simplifyWorklist, "SimplifyWorklist");
    print_vreg_set(freezeWorklist, "FreezeWorklist");
    print_vreg_set(spillWorklist, "SpillWorklist");
    print_vreg_set(coalescedNodes, "CoalescedNodes");
    print_vreg_set(spilledNodes, "SpilledNodes");

    std::cerr << "SelectStack (" << selectStack.size() << "): { ";
    for(unsigned v : selectStack) std::cerr << "%vreg" << v << " ";
    std::cerr << "}\n";

    std::cerr << "WorklistMoves (" << worklistMoves.size() << ")\n";
    std::cerr << "ActiveMoves (" << activeMoves.size() << ")\n";

    size_t final_nodes = coalescedNodes.size() + spilledNodes.size() + selectStack.size();
    std::cerr << "Total Final Nodes: " << final_nodes << "\n";
    std::cerr << "=======================================================\n";
}

std::string RISCv64RegAlloc::regToString(PhysicalReg reg) {
    switch (reg) {
        case PhysicalReg::ZERO: return "x0";  case PhysicalReg::RA: return "ra";
        case PhysicalReg::SP: return "sp";    case PhysicalReg::GP: return "gp";
        case PhysicalReg::TP: return "tp";    case PhysicalReg::T0: return "t0";
        case PhysicalReg::T1: return "t1";    case PhysicalReg::T2: return "t2";
        case PhysicalReg::S0: return "s0";    case PhysicalReg::S1: return "s1";
        case PhysicalReg::A0: return "a0";    case PhysicalReg::A1: return "a1";
        case PhysicalReg::A2: return "a2";    case PhysicalReg::A3: return "a3";
        case PhysicalReg::A4: return "a4";    case PhysicalReg::A5: return "a5";
        case PhysicalReg::A6: return "a6";    case PhysicalReg::A7: return "a7";
        case PhysicalReg::S2: return "s2";    case PhysicalReg::S3: return "s3";
        case PhysicalReg::S4: return "s4";    case PhysicalReg::S5: return "s5";
        case PhysicalReg::S6: return "s6";    case PhysicalReg::S7: return "s7";
        case PhysicalReg::S8: return "s8";    case PhysicalReg::S9: return "s9";
        case PhysicalReg::S10: return "s10";  case PhysicalReg::S11: return "s11";
        case PhysicalReg::T3: return "t3";    case PhysicalReg::T4: return "t4";
        case PhysicalReg::T5: return "t5";    case PhysicalReg::T6: return "t6";
        case PhysicalReg::F0: return "f0";    case PhysicalReg::F1: return "f1";
        case PhysicalReg::F2: return "f2";    case PhysicalReg::F3: return "f3";
        case PhysicalReg::F4: return "f4";    case PhysicalReg::F5: return "f5";
        case PhysicalReg::F6: return "f6";    case PhysicalReg::F7: return "f7";
        case PhysicalReg::F8: return "f8";    case PhysicalReg::F9: return "f9";
        case PhysicalReg::F10: return "f10";  case PhysicalReg::F11: return "f11";
        case PhysicalReg::F12: return "f12";  case PhysicalReg::F13: return "f13";
        case PhysicalReg::F14: return "f14";  case PhysicalReg::F15: return "f15";
        case PhysicalReg::F16: return "f16";  case PhysicalReg::F17: return "f17";
        case PhysicalReg::F18: return "f18";  case PhysicalReg::F19: return "f19";
        case PhysicalReg::F20: return "f20";  case PhysicalReg::F21: return "f21";
        case PhysicalReg::F22: return "f22";  case PhysicalReg::F23: return "f23";
        case PhysicalReg::F24: return "f24";  case PhysicalReg::F25: return "f25";
        case PhysicalReg::F26: return "f26";  case PhysicalReg::F27: return "f27";
        case PhysicalReg::F28: return "f28";  case PhysicalReg::F29: return "f29";
        case PhysicalReg::F30: return "f30";  case PhysicalReg::F31: return "f31";
        default: return "UNKNOWN_REG";
    }
}

std::string RISCv64RegAlloc::regIdToString(unsigned id) {
    const unsigned offset = static_cast<unsigned>(PhysicalReg::PHYS_REG_START_ID);

    // 使用更健壮的检查方式
    if (id >= offset && precolored.count(id)) {
        // 先减去偏移量，得到原始的、小的枚举值
        PhysicalReg reg = static_cast<PhysicalReg>(id - offset);
        // 再将原始枚举值传给 regToString
        return regToString(reg);
    } else {
        return "%vreg" + std::to_string(id);
    }
}

} // namespace sysy