nudt-compiler-cpp/include/mir/MIR.h

#pragma once

#include <initializer_list>
#include <iosfwd>
#include <memory>
#include <string>
#include <vector>

namespace ir {
class Module;
}

namespace mir {

class MIRContext {
 public:
  MIRContext() = default;
};

MIRContext& DefaultContext();

enum class PhysReg {
  W0, W1, W2, W3, W4, W5, W6, W7,
  W8, W9, W10, W11,
  W19, W20, W21, W22, W23, W24,
  X0, X1, X2, X3, X4, X5, X6, X7,
  X8, X9, X10, X11, X29, X30, SP,
  X19, X20, X21, X22, X23, X24,
  S0, S1, S2, S3, S4, S5, S6, S7,  // 单精度浮点寄存器
  S8, S9, S10
};

const char* PhysRegName(PhysReg reg);

enum class Opcode {
  Prologue,
  Epilogue,
  MovImm,
  MovReg,
  FMovImm,     // 浮点立即数加载
  FMovReg,     // 浮点寄存器移动
  LoadStack,
  StoreStack,
  LoadStackOffset,   // 加载数组元素：ldr w8, [x29, base_offset + element_offset]
  StoreStackOffset,  // 存储数组元素：str w8, [x29, base_offset + element_offset]
  LoadStackAddr,     // 加载栈地址：add x9, x29, #offset（用于数组基址）
  LoadIndirect,      // 间接加载：ldr w8, [x9]
  StoreIndirect,     // 间接存储：str w8, [x9]
  LoadIndirectScaled,   // 间接加载：ldr w8, [x9, w10, uxtw #2]
  StoreIndirectScaled,  // 间接存储：str w8, [x9, w10, uxtw #2]
  LoadGlobal,
  StoreGlobal,
  LoadGlobalAddr,  // 加载全局变量地址（用于数组）
  AddRI,
  SubRI,
  AddRR,
  AddRR_UXTW,  // add xN, xM, wK, uxtw（零扩展W寄存器后加到X寄存器）
  SubRR,
  MulRR,
  DivRR,
  ModRR,
  LsrRI,
  LslRI,
  LslRR,   // 逻辑左移（用于 index * 4）
  FAddRR,  // 浮点加法
  FSubRR,  // 浮点减法
  FMulRR,  // 浮点乘法
  FDivRR,  // 浮点除法
  FSqrtRR, // 浮点平方根
  SIToFP,  // 有符号整型转浮点
  FPToSI,  // 浮点转有符号整型
  CmpOnlyRR,
  FCmpOnlyRR,
  CmpRR,
  FCmpRR,  // 浮点比较
  Bl,
  B,       // 无条件跳转
  Bcond,   // 条件跳转（基于之前的 cmp）
  FBcond,  // 浮点条件跳转（基于之前的 fcmp，使用 IEEE 754 兼容的条件码）
  Cbnz,    // 非零跳转
  Cbz,     // 零跳转
  Ret,
};

// 虚拟寄存器类别
enum class VRegClass {
  GPR,    // 通用寄存器 (w0-w11)
  GPR64,  // 64位通用寄存器 (x0-x11)
  FPR,    // 浮点寄存器 (s0-s10)
};

class Operand {
 public:
  enum class Kind { Reg, VReg, Imm, FrameIndex, Symbol };

  static Operand Reg(PhysReg reg);
  static Operand VReg(int vreg_id, VRegClass rc = VRegClass::GPR);
  static Operand Imm(int value);
  static Operand FrameIndex(int index);
  static Operand Symbol(std::string name);

  Kind GetKind() const { return kind_; }
  bool IsReg() const { return kind_ == Kind::Reg; }
  bool IsVReg() const { return kind_ == Kind::VReg; }
  bool IsImm() const { return kind_ == Kind::Imm; }
  bool IsFrameIndex() const { return kind_ == Kind::FrameIndex; }
  bool IsSymbol() const { return kind_ == Kind::Symbol; }

  PhysReg GetReg() const { return reg_; }
  int GetVRegId() const { return vreg_id_; }
  VRegClass GetVRegClass() const { return vreg_class_; }
  int GetImm() const { return imm_; }
  int GetFrameIndex() const { return imm_; }
  const std::string& GetSymbol() const { return symbol_; }

  // 用于寄存器分配后替换 VReg → Reg
  void AssignPhysReg(PhysReg reg);

 private:
  Operand(Kind kind, PhysReg reg, int imm, std::string symbol = "");
  Operand(int vreg_id, VRegClass rc);

  Kind kind_;
  PhysReg reg_ = PhysReg::W0;
  int vreg_id_ = -1;
  VRegClass vreg_class_ = VRegClass::GPR;
  int imm_ = 0;
  std::string symbol_;
};

class MachineInstr {
 public:
  MachineInstr(Opcode opcode, std::vector<Operand> operands = {});

  Opcode GetOpcode() const { return opcode_; }
  const std::vector<Operand>& GetOperands() const { return operands_; }
  std::vector<Operand>& GetOperands() { return operands_; }

  void SetOperand(size_t i, Operand op);

 private:
  Opcode opcode_;
  std::vector<Operand> operands_;
};

struct FrameSlot {
  int index = 0;
  int size = 4;
  int offset = 0;
};

class MachineBasicBlock {
 public:
  explicit MachineBasicBlock(std::string name);

  const std::string& GetName() const { return name_; }
  std::vector<MachineInstr>& GetInstructions() { return instructions_; }
  const std::vector<MachineInstr>& GetInstructions() const { return instructions_; }

  MachineInstr& Append(Opcode opcode,
                       std::initializer_list<Operand> operands = {});

  // CFG 支持
  void AddSuccessor(MachineBasicBlock* succ);
  void AddPredecessor(MachineBasicBlock* pred);
  const std::vector<MachineBasicBlock*>& GetSuccessors() const { return successors_; }
  const std::vector<MachineBasicBlock*>& GetPredecessors() const { return predecessors_; }
  void ClearCFG() { successors_.clear(); predecessors_.clear(); }

 private:
  std::string name_;
  std::vector<MachineInstr> instructions_;
  std::vector<MachineBasicBlock*> successors_;
  std::vector<MachineBasicBlock*> predecessors_;
};

class MachineFunction {
 public:
  explicit MachineFunction(std::string name);

  const std::string& GetName() const { return name_; }
  MachineBasicBlock& GetEntry() { return *blocks_.front(); }
  const MachineBasicBlock& GetEntry() const { return *blocks_.front(); }

  MachineBasicBlock* CreateBlock(std::string name);
  MachineBasicBlock* FindBlock(const std::string& name);
  const std::vector<std::unique_ptr<MachineBasicBlock>>& GetBlocks() const {
    return blocks_;
  }

  int CreateFrameIndex(int size = 4);
  FrameSlot& GetFrameSlot(int index);
  const FrameSlot& GetFrameSlot(int index) const;
  const std::vector<FrameSlot>& GetFrameSlots() const { return frame_slots_; }
  std::vector<FrameSlot>& GetMutableFrameSlots() { return frame_slots_; }

  int GetFrameSize() const { return frame_size_; }
  void SetFrameSize(int size) { frame_size_ = size; }

  // 虚拟寄存器管理
  int CreateVReg(VRegClass rc = VRegClass::GPR);
  int GetNumVRegs() const { return next_vreg_id_; }

  // Callee-saved 寄存器管理
  void AddUsedCalleeSaved(PhysReg reg);
  const std::vector<PhysReg>& GetUsedCalleeSaved() const { return used_callee_saved_; }

 private:
  std::string name_;
  std::vector<std::unique_ptr<MachineBasicBlock>> blocks_;
  std::vector<FrameSlot> frame_slots_;
  int frame_size_ = 0;
  int next_vreg_id_ = 0;
  std::vector<PhysReg> used_callee_saved_;
};

class MachineModule {
 public:
  MachineModule() = default;
  MachineFunction* CreateFunction(std::string name);
  const std::vector<std::unique_ptr<MachineFunction>>& GetFunctions() const {
    return functions_;
  }

  void AddGlobalVar(std::string name, int init_val, int count, bool is_float,
                    std::vector<int> init_elems = {});
  const std::vector<std::tuple<std::string, int, int, bool, std::vector<int>>>&
  GetGlobalVars() const {
    return global_vars_;
  }

 private:
  std::vector<std::unique_ptr<MachineFunction>> functions_;
  std::vector<std::tuple<std::string, int, int, bool, std::vector<int>>>
      global_vars_;  // (name, init, count, is_float, init_elements)
};

std::unique_ptr<MachineModule> LowerToMIR(const ir::Module& module);
void RunPeephole(MachineFunction& function);
void RunRegAlloc(MachineFunction& function);
void RunLoopSlotPromotion(MachineFunction& function);
void RunFrameLowering(MachineFunction& function);
void PrintAsm(const MachineModule& module, std::ostream& os);

}  // namespace mir