CompilerSysY/include/llir_instruction.h

#pragma once
#include "common.h"
#include "llir_type.h"
#include "llir_value.h"
#include <cassert>
#include <sstream>
#include <string>

namespace CompSysY {

DEF_PTR_T(InstAlloca);
DEF_PTR_T(InstStore);
DEF_PTR_T(InstLoad);
DEF_PTR_T(InstBinary);
DEF_PTR_T(InstZext);
DEF_PTR_T(InstBranch);
DEF_PTR_T(InstReturn);
DEF_PTR_T(InstCall);
DEF_PTR_T(InstGEP);
DEF_PTR_T(InstPhi);

enum class InstTag {
  Add,
  Sub,
  Mod,
  Mul,
  Div,
  Lt,
  Le,
  Ge,
  Gt,
  Eq,
  Ne,
  // And, Or,
  Br,
  Call,
  Ret,
  Alloca,
  Load,
  Store,
  GEP,
  Zext,
  Phi,
  MemPhi,
  LoadDep,
  InsertEle,
  ExtractEle
};

class Instruction : public User {
public:
  int ir_seqno = -1;
  InstTag tag;
  BasicBlockPtr_t parent_bb;
  std::list<InstructionPtr_t>::iterator inst_itr;
  Instruction(InstTag inst_tag, TypePtr_t type, BasicBlockPtr_t parent_bb)
    : User("", type), tag(inst_tag), parent_bb(parent_bb) {}
  std::string tag_string() {
    switch (tag) {
      case InstTag::Add: return "add";
      case InstTag::Sub: return "sub";
      case InstTag::Mod: return "srem";
      case InstTag::Mul: return "mul";
      case InstTag::Div: return "sdiv";
      case InstTag::Lt: return "icmp slt";
      case InstTag::Le: return "icmp sle";
      case InstTag::Ge: return "icmp sge";
      case InstTag::Gt: return "icmp sgt";
      case InstTag::Eq: return "icmp eq";
      case InstTag::Ne: return "icmp ne";
      // case InstTag::And: return "and";
      // case InstTag::Or: return "or";
      case InstTag::Br: return "br";
      case InstTag::Call: return "call";
      case InstTag::Ret: return "ret";
      case InstTag::Alloca: return "alloca";
      case InstTag::Load: return "load";
      case InstTag::Store: return "store";
      case InstTag::GEP: return "GEP";
      case InstTag::Zext: return "zext";
      case InstTag::Phi: return "Phi";
      case InstTag::MemPhi: return "MemPhi";
      case InstTag::LoadDep: return "LoadDep";
      case InstTag::InsertEle: return "InsertEle";
      case InstTag::ExtractEle: return "ExtractEle";
    }
  }
  virtual std::string to_string() override {
    return type->to_string() + " " + tag_string();
  };
};

class InstAlloca : public Instruction {
public:
  InstAlloca(TypePtr_t alloc_type, std::shared_ptr<BasicBlock> parent_bb)
    : Instruction(InstTag::Alloca, std::make_shared<PointerType>(alloc_type), parent_bb) {}
  virtual std::string to_IR_string() override {
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static sptr(InstAlloca) New(const std::string &name, TypePtr_t type, sptr(BasicBlock) parent_bb) {
    auto inst         = std::make_shared<InstAlloca>(type, parent_bb);
    inst->name        = name;
    auto func_head_bb = parent_bb->parent->bb_list.front();
    // put all alloca in the head of each basic block
    inst->inst_itr = func_head_bb->inst_list.insert(func_head_bb->inst_list.begin(), inst);
    return inst;
  }
};

class InstStore : public Instruction {
public:
  InstStore(std::shared_ptr<Value> value, std::shared_ptr<Value> pointer, std::shared_ptr<BasicBlock> parent_bb)
    : Instruction(InstTag::Store, TypeHelper::TYPE_VOID, parent_bb) {
    assert(value);
    Add_Operand(value);
    Add_Operand(pointer);
  }
  static sptr(InstStore) New(sptr(Value) value, sptr(Value) pointer, sptr(BasicBlock) parent_bb) {
    auto inst      = std::make_shared<InstStore>(value, pointer, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

class InstLoad : public Instruction {
public:
  InstLoad(std::shared_ptr<Value> value, TypePtr_t type, std::shared_ptr<BasicBlock> parent_bb)
    : Instruction(InstTag::Load, type, parent_bb) {
    Add_Operand(value);
  }
  virtual std::string to_IR_string() override {
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static sptr(InstLoad) New(std::shared_ptr<Value> value, TypePtr_t type, std::shared_ptr<BasicBlock> parent_bb) {
    auto inst      = std::make_shared<InstLoad>(value, type, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

class InstBinary : public Instruction {
public:
  InstBinary(
    InstTag inst_tag,
    TypePtr_t val_type,
    std::shared_ptr<Value> op1,
    std::shared_ptr<Value> op2,
    std::shared_ptr<BasicBlock> parent_bb
  )
    : Instruction(inst_tag, val_type, parent_bb) {
    Add_Operand(op1);
    Add_Operand(op2);
  }
  virtual std::string to_IR_string() override {
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static sptr(InstBinary) New(
    InstTag inst_tag,
    std::shared_ptr<Value> op1,
    std::shared_ptr<Value> op2,
    std::shared_ptr<BasicBlock> parent_bb
  ) {
    std::shared_ptr<InstBinary> inst;
    if (InstTag::Lt <= inst_tag && inst_tag <= InstTag::Ne) {
      inst = std::make_shared<InstBinary>(inst_tag, TypeHelper::TYPE_I1, op1, op2, parent_bb);
    }
    else if (InstTag::Add <= inst_tag && inst_tag <= InstTag::Div) {
      inst = std::make_shared<InstBinary>(inst_tag, TypeHelper::TYPE_I32, op1, op2, parent_bb);
    }
    else {
      panic("Invalid Binary Operation");
    }
    assert(Type::isType<IntegerType>(op1->type));
    assert(Type::isType<IntegerType>(op2->type));
    if (!(Type::asType<IntegerType>(op1->type)->int_type == Type::asType<IntegerType>(op2->type)->int_type)) {
      if (!(shared_cast<ConstantInt>(op1) && shared_cast<ConstantInt>(op1)->value == 0)
      && !(shared_cast<ConstantInt>(op2) && shared_cast<ConstantInt>(op2)->value == 0)) {
        LOG(ERROR) << op1->to_string() << ",\t" << op2->to_string();
        assert(0);
      }
    }
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

//<result> = zext <ty> <value> to <ty2>
// zext is integer type, i32
class InstZext : public Instruction {
public:
  InstZext(std::shared_ptr<Value> op, std::shared_ptr<BasicBlock> parent_bb)
    : Instruction(InstTag::Zext, TypeHelper::TYPE_I32, parent_bb) {
    Add_Operand(op);
  }
  virtual std::string to_IR_string() override {
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static sptr(InstZext) New(std::shared_ptr<Value> op, std::shared_ptr<BasicBlock> parent_bb) {
    auto inst      = std::make_shared<InstZext>(op, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

class InstBranch : public Instruction {
public:
  // conditional branch
  InstBranch(ValuePtr_t cond, BasicBlockPtr_t true_block, BasicBlockPtr_t false_block, BasicBlockPtr_t parent_bb)
    : Instruction(InstTag::Br, TypeHelper::TYPE_VOID, parent_bb) {
    this->Add_Operand(cond);
    this->Add_Operand(true_block);
    this->Add_Operand(false_block);
  }
  static sptr(InstBranch)
    New(ValuePtr_t cond, BasicBlockPtr_t true_block, BasicBlockPtr_t false_block, BasicBlockPtr_t parent_bb) {
    auto inst      = std::make_shared<InstBranch>(cond, true_block, false_block, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
  // unconditional branch
  InstBranch(BasicBlockPtr_t target_block, BasicBlockPtr_t parent_bb)
    : Instruction(InstTag::Br, TypeHelper::TYPE_VOID, parent_bb) {
    this->Add_Operand(target_block);
  }
  static sptr(InstBranch) New(BasicBlockPtr_t target_block, BasicBlockPtr_t parent_bb) {
    auto inst      = std::make_shared<InstBranch>(target_block, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

class InstReturn : public Instruction {
public:
  InstReturn(ValuePtr_t ret_val, BasicBlockPtr_t parent_bb)
    : Instruction(InstTag::Ret, TypeHelper::TYPE_VOID, parent_bb) {
    this->Add_Operand(ret_val);
  }
  InstReturn(BasicBlockPtr_t parent_bb) : Instruction(InstTag::Ret, TypeHelper::TYPE_VOID, parent_bb) {}
  static sptr(InstReturn) New(ValuePtr_t ret_val, BasicBlockPtr_t parent_bb) {
    auto inst      = std::make_shared<InstReturn>(ret_val, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
  static sptr(InstReturn) New(BasicBlockPtr_t parent_bb) {
    auto inst      = std::make_shared<InstReturn>(parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

// call's type is the function's return type
class InstCall : public Instruction {
public:
  InstCall(FunctionPtr_t func, const std::vector<ValuePtr_t> &args, BasicBlockPtr_t parent_bb)
    : Instruction(InstTag::Call, func->get_type()->return_type, parent_bb) {
    Add_Operand(func);
    for (auto arg : args) {
      Add_Operand(arg);
    }
  }
  virtual std::string to_IR_string() override {
    if (Type::isType<VoidType>(type)) {
      panic("No ret val");
    }
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static std::shared_ptr<InstCall> New(
    FunctionPtr_t func,
    const std::vector<ValuePtr_t> &args,
    BasicBlockPtr_t parent_bb
  ) {
    auto inst      = std::make_shared<InstCall>(func, args, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

// getelementptr <ty>, ptr <ptrval>{, [inrange] <ty> <idx>}*
// Example: lea arr[2][3] from arr[5][4]:
//  &arr[2][3] = GEP [5x[4xi32]], [5x[4xi32]]* arr, i32 0, i32 2, i32 3
// For simplicity, we want to limit the use of gep of at most 3 ops
// op0: pointer to array;
// if there is only one index, the return value has the same type as op0
// if there is 2 index, the return value has type of arr's element
class InstGEP : public Instruction {
public:
  ValuePtr_t aim_to;
  TypePtr_t element_type;
  InstGEP(ValuePtr_t pointer, const std::vector<ValuePtr_t> &indices, BasicBlockPtr_t parent_bb)
    : Instruction(InstTag::GEP, std::make_shared<PointerType>(extract_type(pointer, indices)), parent_bb) {
    if (shared_cast<InstGEP>(pointer)) {
      aim_to = std::dynamic_pointer_cast<InstGEP>(pointer)->aim_to;
    }
    else if (shared_cast<InstAlloca>(pointer) || shared_cast<GlobalVar>(pointer)) {
      aim_to = pointer;
    }
    else {
      // LOG(WARNING) << "Unexpected pointer type " << pointer->to_string();
      aim_to = nullptr;
    }
    assert(indices.size() <= 2);
    element_type = extract_type(pointer, indices);
    Add_Operand(pointer);
    for (auto index : indices) {
      Add_Operand(index);
    }
  }

  virtual std::string to_IR_string() override {
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static std::shared_ptr<InstGEP> New(
    ValuePtr_t pointer,
    const std::vector<ValuePtr_t> &indices,
    BasicBlockPtr_t parent_bb
  ) {
    auto inst      = std::make_shared<InstGEP>(pointer, indices, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }

  // get the inner
  static TypePtr_t extract_type(ValuePtr_t pointer, const std::vector<ValuePtr_t> &indices) {
    auto pointed_type = shared_cast<PointerType>(pointer->type)->pointed_type;
    if (Type::isType<IntegerType>(pointed_type)) {
      return pointed_type;
    }
    else if (Type::isType<ArrayType>(pointed_type)) {
      for (int i = 1; i < indices.size(); ++i) {
        pointed_type = shared_cast<ArrayType>(pointed_type)->elem_type;
      }
      return pointed_type;
    }
    LOG(ERROR) << "unexpected type: " << pointer->to_string();
    assert(0);
  }
};

class InstPhi : public Instruction {
public:
  InstPhi(TypePtr_t type, const decltype(Function::bb_list) &incoming_vals, BasicBlockPtr_t parent_bb)
    : Instruction(InstTag::Phi, type, parent_bb) {
    operand_list.resize(incoming_vals.size(), nullptr);
  }
  void set_incoming_val(unsigned index, ValuePtr_t val) {
    auto old_op         = operand_list[index];
    operand_list[index] = val;
    if (val) val->use_list.push_back({/*val.get(),*/ this, index});
    if (old_op) {
      if (std::find(operand_list.begin(), operand_list.end(), old_op) == operand_list.end()) {
        old_op->u_remove_use(this);
      }
    }
  }
  virtual std::string to_IR_string() override {
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static InstPhiPtr_t New(
    TypePtr_t type,
    const decltype(Function::bb_list) &incoming_vals,
    BasicBlockPtr_t parent_bb,
    const std::string &name
  ) {
    auto inst      = std::make_shared<InstPhi>(type, incoming_vals, parent_bb);
    inst->name     = name;
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.begin(), inst);
    return inst;
  }
};

} // namespace CompSysY