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();

// AArch64 physical registers
enum class PhysReg { 
  W0, W1, W2, W3, W4, W5, W6, W7,
  W8, W9, W10, W11, W12, W13, W14, W15,
  X0, X1, X2, X3, X4, X5, X6, X7,
  X8, X9, X10, X11, X12, X13, X14, X15,
  X16, X17,
  S0, S1, S2, S3, S4, S5, S6, S7,
  S8, S9, S10, S11, S12, S13, S14, S15,
  X29, X30, SP, WZR, XZR 
};

const char* PhysRegName(PhysReg reg);

enum class Opcode {
  Prologue,
  Epilogue,
  MovImm,
  MovRR,
  LoadStack,
  StoreStack,
  AddrStack,
  LoadGlobal,
  StoreGlobal,
  AddRR,
  AddRRI,
  AddRRR_LSL,
  SubRR,
  MulRR,
  SDivRR,
  MSubRRR,
  Sxtw,
  NegR,
  CmpRR,
  CSet,
  FAdd,
  FSub,
  FMUL,
  FDiv,
  FNeg,
  FCmp,
  FCvtSI2FP,
  FCvtFP2SI,
  LoadR,
  StoreR,
  Call,
  B,
  BCond,
  Ret,
};

enum class CondCode { EQ, NE, LT, LE, GT, GE };

class Operand {
 public:
  enum class Kind { Reg, Imm, FrameIndex, Label, Global, Cond };

  static Operand Reg(PhysReg reg);
  static Operand Imm(int value);
  static Operand FrameIndex(int index);
  static Operand Label(const std::string& name);
  static Operand Global(const std::string& name);
  static Operand Cond(CondCode cc);

  Kind GetKind() const { return kind_; }
  PhysReg GetReg() const { return reg_; }
  int GetImm() const { return imm_; }
  int GetFrameIndex() const { return imm_; }
  const std::string& GetLabel() const { return label_; }
  const std::string& GetGlobal() const { return label_; }
  CondCode GetCond() const { return static_cast<CondCode>(imm_); }

 private:
  Operand(Kind kind, PhysReg reg, int imm, std::string label = "");

  Kind kind_;
  PhysReg reg_;
  int imm_;
  std::string label_;
};

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

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

 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 = {});

 private:
  std::string name_;
  std::vector<MachineInstr> instructions_;
};

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

  const std::string& GetName() const { return name_; }
  
  MachineBasicBlock& CreateBlock(const std::string& name);
  std::vector<std::unique_ptr<MachineBasicBlock>>& GetBlocks() { return blocks_; }
  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_; }

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

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

struct GlobalVariable {
  std::string name;
  int init_value = 0;
  size_t size = 4;
  bool is_const = false;
};

class MachineModule {
 public:
  MachineModule() = default;
  std::vector<std::unique_ptr<MachineFunction>>& GetFunctions() { return functions_; }
  const std::vector<std::unique_ptr<MachineFunction>>& GetFunctions() const { return functions_; }
  
  std::vector<GlobalVariable>& GetGlobals() { return globals_; }
  const std::vector<GlobalVariable>& GetGlobals() const { return globals_; }

 private:
  std::vector<std::unique_ptr<MachineFunction>> functions_;
  std::vector<GlobalVariable> globals_;
};

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

}  // namespace mir