#include "RISCv64Passes.h"
#include <iostream>

namespace sysy {

// --- 寄存器分配前优化 ---

void PreRA_Scheduler::runOnMachineFunction(MachineFunction* mfunc) {
    // TODO: 在此实现寄存器分配前的指令调度。
    // 遍历mfunc中的每一个MachineBasicBlock。
    // 对每个基本块内的MachineInstr列表进行重排。
    //
    // 实现思路:
    // 1. 分析每个基本块内指令的数据依赖关系，构建依赖图(DAG)。
    // 2. 根据目标处理器的流水线特性（指令延迟等），使用列表调度等算法对指令进行重排。
    // 3. 此时操作的是虚拟寄存器，只存在真依赖，调度自由度最大。
    //
    // std::cout << "Running Pre-RA Instruction Scheduler..." << std::endl;
}


// --- 寄存器分配后优化 ---

void PeepholeOptimizer::runOnMachineFunction(MachineFunction* mfunc) {
    // TODO: 在此实现窥孔优化。
    // 遍历mfunc中的每一个MachineBasicBlock。
    // 对每个基本块内的MachineInstr列表进行扫描和替换。
    //
    // 实现思路:
    // 1. 维护一个大小固定（例如3-5条指令）的滑动窗口。
    // 2. 识别特定的冗余模式，例如:
    //    - `mv a0, a1` 后紧跟 `mv a1, a0` (可消除的交换)
    //    - `sw t0, 12(s0)` 后紧跟 `lw t1, 12(s0)` (冗余加载)
    //    - 强度削减: `mul x, x, 2` -> `slli x, x, 1`
    // 3. 识别后，直接修改MachineInstr列表（删除、替换或插入指令）。
    //
    // std::cout << "Running Post-RA Peephole Optimizer..." << std::endl;
}

void PostRA_Scheduler::runOnMachineFunction(MachineFunction* mfunc) {
    // TODO: 在此实现寄存器分配后的局部指令调度。
    // 遍历mfunc中的每一个MachineBasicBlock。
    // 重点关注由寄存器分配器插入的spill/fill代码。
    //
    // 实现思路:
    // 1. 识别出用于spill/fill的lw/sw指令。
    // 2. 在不违反数据依赖（包括物理寄存器引入的伪依赖）的前提下，
    //    尝试将lw指令向上移动，使其与使用它的指令之间有足够的距离，以隐藏访存延迟。
    // 3. 同样，可以尝试将sw指令向下移动。
    //
    // std::cout << "Running Post-RA Local Scheduler..." << std::endl;
}

void CalleeSavedHandler::runOnMachineFunction(MachineFunction* mfunc) {
    StackFrameInfo& frame_info = mfunc->getFrameInfo();
    std::set<PhysicalReg> used_callee_saved;

    // 1. 扫描所有指令，找出被使用的s寄存器
    for (auto& mbb : mfunc->getBlocks()) {
        for (auto& instr : mbb->getInstructions()) {
            for (auto& op : instr->getOperands()) {
                
                // 辅助Lambda，用于检查和插入寄存器
                auto check_and_insert_reg = [&](RegOperand* reg_op) {
                    if (!reg_op->isVirtual()) {
                        PhysicalReg preg = reg_op->getPReg();
                        // --- 关键检查点 ---
                        // 必须严格判断是否在 s0-s11 的范围内。
                        // a0, t0 等寄存器绝对不应被视为被调用者保存寄存器。
                        if (preg >= PhysicalReg::S0 && preg <= PhysicalReg::S11) {
                            used_callee_saved.insert(preg);
                        }
                    }
                };

                if (op->getKind() == MachineOperand::KIND_REG) {
                    check_and_insert_reg(static_cast<RegOperand*>(op.get()));
                } else if (op->getKind() == MachineOperand::KIND_MEM) {
                    check_and_insert_reg(static_cast<MemOperand*>(op.get())->getBase());
                }
            }
        }
    }
    
    // 如果没有使用s寄存器（除了可能作为帧指针的s0），则无需操作
    if (used_callee_saved.empty() || (used_callee_saved.size() == 1 && used_callee_saved.count(PhysicalReg::S0))) {
        return;
    }

    // 将结果存入StackFrameInfo，供后续使用
    frame_info.used_callee_saved_regs = used_callee_saved;

    // 2. 在函数序言插入保存指令
    MachineBasicBlock* entry_block = mfunc->getBlocks().front().get();
    auto& entry_instrs = entry_block->getInstructions();
    auto prologue_end = entry_instrs.begin();
    
    // 找到序言结束的位置（通常是addi s0, sp, size之后）
    for (auto it = entry_instrs.begin(); it != entry_instrs.end(); ++it) {
        if ((*it)->getOpcode() == RVOpcodes::ADDI && 
            (*it)->getOperands()[0]->getKind() == MachineOperand::KIND_REG &&
            static_cast<RegOperand*>((*it)->getOperands()[0].get())->getPReg() == PhysicalReg::S0) 
        {
            prologue_end = std::next(it);
            break;
        }
    }
    
    // 为了栈帧布局确定性，对寄存器进行排序
    std::vector<PhysicalReg> sorted_regs(used_callee_saved.begin(), used_callee_saved.end());
    std::sort(sorted_regs.begin(), sorted_regs.end());
    
    int current_offset = -16; // ra和s0已经占用了-8和-16的位置
    for (PhysicalReg reg : sorted_regs) {
        if (reg == PhysicalReg::S0) continue; // s0已经在序言中处理
        current_offset -= 8;
        auto sd = std::make_unique<MachineInstr>(RVOpcodes::SD);
        sd->addOperand(std::make_unique<RegOperand>(reg));
        sd->addOperand(std::make_unique<MemOperand>(
            std::make_unique<RegOperand>(PhysicalReg::S0), // 假设s0是帧指针
            std::make_unique<ImmOperand>(current_offset)
        ));
        entry_instrs.insert(prologue_end, std::move(sd));
    }

    // 3. 在函数结尾（ret之前）插入恢复指令
    for (auto& mbb : mfunc->getBlocks()) {
        for (auto it = mbb->getInstructions().begin(); it != mbb->getInstructions().end(); ++it) {
            if ((*it)->getOpcode() == RVOpcodes::RET) {
                // 以相反的顺序恢复
                current_offset = -16;
                for (PhysicalReg reg : sorted_regs) {
                    if (reg == PhysicalReg::S0) continue;
                    current_offset -= 8;
                    auto ld = std::make_unique<MachineInstr>(RVOpcodes::LD);
                    ld->addOperand(std::make_unique<RegOperand>(reg));
                    ld->addOperand(std::make_unique<MemOperand>(
                        std::make_unique<RegOperand>(PhysicalReg::S0),
                        std::make_unique<ImmOperand>(current_offset)
                    ));
                    mbb->getInstructions().insert(it, std::move(ld));
                }
                break; // 处理完一个基本块的ret即可
            }
        }
    }
}

} // namespace sysy