#include "EliminateFrameIndices.h"
#include "RISCv64ISel.h"
#include <cassert>
#include <vector> // [新增] 为插入指令而包含

namespace sysy {

// getTypeSizeInBytes 是一个通用辅助函数，保持不变
unsigned EliminateFrameIndicesPass::getTypeSizeInBytes(Type* type) {
    if (!type) {
        assert(false && "Cannot get size of a null type.");
        return 0;
    }

    switch (type->getKind()) {
        case Type::kInt:
        case Type::kFloat:
            return 4;
        case Type::kPointer:
            return 8;
        case Type::kArray: {
            auto arrayType = type->as<ArrayType>();
            return arrayType->getNumElements() * getTypeSizeInBytes(arrayType->getElementType());
        }
        default:
            assert(false && "Unsupported type for size calculation.");
            return 0;
    }
}

void EliminateFrameIndicesPass::runOnMachineFunction(MachineFunction* mfunc) {
    StackFrameInfo& frame_info = mfunc->getFrameInfo();
    Function* F = mfunc->getFunc();
    RISCv64ISel* isel = mfunc->getISel();
    
    // 1. [已移除] 不再处理栈传递的参数
    // 原先处理栈参数 (arg_idx >= 8) 的逻辑已被移除。
    // 这项职责已完全转移到 PrologueEpilogueInsertionPass，以避免逻辑冲突和错误。
    // [注释更新] -> 上述注释已过时。根据新方案，我们将在这里处理栈传递的参数，
    // 以便在寄存器分配前就将数据流显式化，修复溢出逻辑的BUG。

    // 2. 只为局部变量(AllocaInst)分配栈空间和计算偏移量
    // 局部变量从 s0 下方（负偏移量）开始分配，紧接着为 ra 和 s0 预留的16字节之后
    int local_var_offset = 16;
    
    if(F) { // 确保函数指针有效
        for (auto& bb : F->getBasicBlocks()) {
            for (auto& inst : bb->getInstructions()) {
                if (auto alloca = dynamic_cast<AllocaInst*>(inst.get())) {
                    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字节

                    local_var_offset += size;
                    unsigned alloca_vreg = isel->getVReg(alloca);
                    // 局部变量使用相对于s0的负向偏移
                    frame_info.alloca_offsets[alloca_vreg] = -local_var_offset;
                }
            }
        }
    }
    
    // 记录仅由AllocaInst分配的局部变量的总大小
    frame_info.locals_size = local_var_offset - 16;

    // 3. [核心修改] 在函数入口为所有栈传递的参数插入load指令
    // 这个步骤至关重要：它在寄存器分配之前，为这些参数的vreg创建了明确的“定义(def)”指令。
    // 这解决了在高寄存器压力下，当这些vreg被溢出时，`rewriteProgram`找不到其定义点而崩溃的问题。
    if (F && isel && !mfunc->getBlocks().empty()) {
        MachineBasicBlock* entry_block = mfunc->getBlocks().front().get();
        std::vector<std::unique_ptr<MachineInstr>> arg_load_instrs;
        int arg_idx = 0;
        for (Argument* arg : F->getArguments()) {
            // 根据ABI，前8个整型/指针参数通过寄存器传递，这里只处理超出部分。
            if (arg_idx >= 8) {
                // 计算参数在调用者栈帧中的位置，该位置相对于被调用者的帧指针s0是正向偏移。
                // 第9个参数(arg_idx=8)位于 0(s0)，第10个(arg_idx=9)位于 8(s0)，以此类推。
                int offset = (arg_idx - 8) * 8; 
                unsigned arg_vreg = isel->getVReg(arg);
                Type* arg_type = arg->getType();

                // 根据参数类型选择正确的加载指令
                RVOpcodes load_op;
                if (arg_type->isFloat()) {
                    load_op = RVOpcodes::FLW; // 单精度浮点
                } else if (arg_type->isPointer()) {
                    load_op = RVOpcodes::LD;  // 64位指针
                } else {
                    load_op = RVOpcodes::LW;  // 32位整数
                }
                
                // 创建加载指令: lw/ld/flw vreg, offset(s0)
                auto load_instr = std::make_unique<MachineInstr>(load_op);
                load_instr->addOperand(std::make_unique<RegOperand>(arg_vreg));
                load_instr->addOperand(std::make_unique<MemOperand>(
                    std::make_unique<RegOperand>(PhysicalReg::S0), // 基址为帧指针
                    std::make_unique<ImmOperand>(offset)
                ));
                arg_load_instrs.push_back(std::move(load_instr));
            }
            arg_idx++;
        }
        
        // 将所有新创建的参数加载指令一次性插入到入口块的起始位置
        auto& entry_instrs = entry_block->getInstructions();
        entry_instrs.insert(entry_instrs.begin(),
                            std::make_move_iterator(arg_load_instrs.begin()),
                            std::make_move_iterator(arg_load_instrs.end()));
    }


    // 4. 遍历所有机器指令，将访问局部变量的伪指令展开为真实指令
    // 由于处理参数的逻辑已移除，这里的展开现在只针对局部变量，因此是正确的。
    // [注释更新] -> 上述注释已过时。此部分逻辑保持不变，它正确地处理了局部变量。
    for (auto& mbb : mfunc->getBlocks()) {
        std::vector<std::unique_ptr<MachineInstr>> new_instructions;
        for (auto& instr_ptr : mbb->getInstructions()) {
            RVOpcodes opcode = instr_ptr->getOpcode();

            if (opcode == RVOpcodes::FRAME_LOAD_W || opcode == RVOpcodes::FRAME_LOAD_D || opcode == RVOpcodes::FRAME_LOAD_F) {
                RVOpcodes real_load_op;
                if (opcode == RVOpcodes::FRAME_LOAD_W) real_load_op = RVOpcodes::LW;
                else if (opcode == RVOpcodes::FRAME_LOAD_D) real_load_op = RVOpcodes::LD;
                else real_load_op = RVOpcodes::FLW;

                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(Type::getPointerType(Type::getIntType()));

                // 展开为: 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/flw 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 || opcode == RVOpcodes::FRAME_STORE_F) {
                RVOpcodes real_store_op;
                if (opcode == RVOpcodes::FRAME_STORE_W) real_store_op = RVOpcodes::SW;
                else if (opcode == RVOpcodes::FRAME_STORE_D) real_store_op = RVOpcodes::SD;
                else real_store_op = RVOpcodes::FSW;

                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(Type::getPointerType(Type::getIntType()));
                
                // 展开为: 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/fsw 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));
            }
        }
        mbb->getInstructions() = std::move(new_instructions);
    }
}

} // namespace sysy