nudt-compiler-cpp/src/ir/IR.h

// 极简 IR 定义：当前只支撑 i32 和加法，演示用。
// 可在此基础上扩展更多类型/指令
#pragma once

#include <memory>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

namespace ir {

class Type;
class ConstantInt;
class Instruction;
class BasicBlock;

// IR 上下文：集中管理类型、常量等共享资源，便于复用与扩展。
class Context {
 public:
  ~Context();
  const std::shared_ptr<Type>& Void();
  const std::shared_ptr<Type>& Int32();
  const std::shared_ptr<Type>& PtrInt32();
  // 去重创建 i32 常量。
  ConstantInt* GetConstInt(int v);
 
  std::string NextTemp();

 private:
  std::shared_ptr<Type> void_;
  std::shared_ptr<Type> int32_;
  std::shared_ptr<Type> ptr_i32_;
  std::unordered_map<int, std::unique_ptr<ConstantInt>> const_ints_;
  int temp_index_ = -1;
};

Context& DefaultContext();

class Type {
 public:
  enum class Kind { Void, Int32, PtrInt32 };
  explicit Type(Kind k) : kind_(k) {}
  Kind kind() const { return kind_; }
  static std::shared_ptr<Type> Void();
  static std::shared_ptr<Type> Int32();
  static std::shared_ptr<Type> PtrInt32();

 private:
  Kind kind_;
};

class Value {
 public:
  Value(std::shared_ptr<Type> ty, std::string name)
      : type_(std::move(ty)), name_(std::move(name)) {}
  virtual ~Value() = default;
  const std::shared_ptr<Type>& type() const { return type_; }
  const std::string& name() const { return name_; }
  void set_name(std::string n) { name_ = std::move(n); }
  void AddUser(Instruction* user) { users_.push_back(user); }
  const std::vector<Instruction*>& users() const { return users_; }

 protected:
  std::shared_ptr<Type> type_;
  std::string name_;
  std::vector<Instruction*> users_;
};

class ConstantInt : public Value {
 public:
  explicit ConstantInt(int v);
  int value() const { return value_; }

 private:
  int value_{};
};

enum class Opcode { Add, Sub, Mul, Alloca, Load, Store, Ret };

class Instruction : public Value {
 public:
  Instruction(Opcode op, std::shared_ptr<Type> ty, std::string name = "")
      : Value(std::move(ty), std::move(name)), opcode_(op) {}
  Opcode opcode() const { return opcode_; }
  bool IsTerminator() const { return opcode_ == Opcode::Ret; }
  BasicBlock* parent() const { return parent_; }
  void set_parent(BasicBlock* parent) { parent_ = parent; }

 private:
  Opcode opcode_;
  BasicBlock* parent_ = nullptr;
};

class BinaryInst : public Instruction {
 public:
  BinaryInst(Opcode op, std::shared_ptr<Type> ty, Value* lhs, Value* rhs,
             std::string name);
  Value* lhs() const { return lhs_; }
  Value* rhs() const { return rhs_; }

 private:
  Value* lhs_;
  Value* rhs_;
};

class ReturnInst : public Instruction {
 public:
  explicit ReturnInst(Value* val);
  Value* value() const { return value_; }

 private:
  Value* value_;
};

class AllocaInst : public Instruction {
 public:
  explicit AllocaInst(std::string name);
};

class LoadInst : public Instruction {
 public:
  LoadInst(Value* ptr, std::string name);
  Value* ptr() const { return ptr_; }

 private:
  Value* ptr_;
};

class StoreInst : public Instruction {
 public:
  StoreInst(Value* val, Value* ptr);
  Value* value() const { return value_; }
  Value* ptr() const { return ptr_; }

 private:
  Value* value_;
  Value* ptr_;
};

class BasicBlock {
 public:
  explicit BasicBlock(std::string name) : name_(std::move(name)) {}
  const std::string& name() const { return name_; }
  bool HasTerminator() const {
    return !instructions_.empty() && instructions_.back()->IsTerminator();
  }
  const std::vector<std::unique_ptr<Instruction>>& instructions() const {
    return instructions_;
  }
  template <typename T, typename... Args>
  T* Append(Args&&... args) {
    if (HasTerminator()) {
      throw std::runtime_error("BasicBlock 已有 terminator，不能继续追加指令: " +
                               name_);
    }
    auto inst = std::make_unique<T>(std::forward<Args>(args)...);
    auto* ptr = inst.get();
    ptr->set_parent(this);
    instructions_.push_back(std::move(inst));
    return ptr;
  }

 private:
  std::string name_;
  std::vector<std::unique_ptr<Instruction>> instructions_;
};

class Function : public Value {
 public:
  // 允许显式指定返回类型，便于后续扩展多种函数签名。
  Function(std::string name, std::shared_ptr<Type> ret_type);
  BasicBlock* CreateBlock(const std::string& name);
  BasicBlock* entry() { return entry_; }
  const BasicBlock* entry() const { return entry_; }
  const std::vector<std::unique_ptr<BasicBlock>>& blocks() const {
    return blocks_;
  }

 private:
  BasicBlock* entry_ = nullptr;
  std::vector<std::unique_ptr<BasicBlock>> blocks_;
};

class Module {
 public:
  // 创建函数时显式传入返回类型，便于在 IRGen 中根据语法树信息选择类型。
  Function* CreateFunction(const std::string& name,
                           std::shared_ptr<Type> ret_type);
  const std::vector<std::unique_ptr<Function>>& functions() const {
    return functions_;
  }

 private:
  std::vector<std::unique_ptr<Function>> functions_;
};

class IRBuilder {
 public:
  explicit IRBuilder(BasicBlock* bb) : insertBlock_(bb) {}
  void SetInsertPoint(BasicBlock* bb) { insertBlock_ = bb; }
  BasicBlock* GetInsertBlock() const { return insertBlock_; }

  // 构造常量、二元运算、返回指令的最小集合。
  ConstantInt* CreateConstInt(int v);
  BinaryInst* CreateBinary(Opcode op, Value* lhs, Value* rhs,
                           const std::string& name);
  BinaryInst* CreateAdd(Value* lhs, Value* rhs, const std::string& name) {
    return CreateBinary(Opcode::Add, lhs, rhs, name);
  }
  AllocaInst* CreateAllocaI32(const std::string& name);
  LoadInst* CreateLoad(Value* ptr, const std::string& name);
  StoreInst* CreateStore(Value* val, Value* ptr);
  ReturnInst* CreateRet(Value* v);

 private:
  BasicBlock* insertBlock_;
};

class IRPrinter {
 public:
  void Print(const Module& module);
};

}  // namespace ir