#pragma once
#include "common.h"
#include "llir.h"

using std::make_shared;

namespace CompSysY {

// a?, t?, ra are caller-saved
// s? are callee-saved
// x0,gp,tp are preserved in user-space
enum class RV64Reg {
  x0 = 0, // zero N/A
  ra,     // ra, caller
  sp,     // sp, callee
  gp,     // gp, N/A
  tp,     // tp, N/A
  t0,     // t0, caller
  t1,
  t2, // t1-2, caller
  s0, // s0/fp, callee
  s1, // s1, callee
  a0,
  a1, // a0-1,caller
  a2,
  a3,
  a4,
  a5,
  a6,
  a7, // a2-7,caller
  s2,
  s3,
  s4,
  s5,
  s6,
  s7,
  s8,
  s9,
  s10,
  s11, // s2-11,callee
  t3,
  t4,
  t5,
  t6, // t3-6,caller
};
inline const char* enum_to_string(const RV64Reg& tag) {
  const static char* _str_tab[] = {"x0","ra","sp","gp","tp","t0","t1","t2","s0","s1","a0","a1","a2","a3","a4","a5","a6","a7","s2","s3","s4","s5","s6","s7","s8","s9","s1","s1","t3","t4","t5","t6",};
  return _str_tab[(unsigned)tag];
}

// riscv calling convention see: https://github.com/riscv-non-isa/riscv-elf-psabi-doc/blob/master/riscv-cc.adoc

const unsigned XLEN      = 8;
const unsigned XLEN_MASK = XLEN - 1;
inline unsigned xlen_rnd_up(unsigned src) {
  return (src & XLEN_MASK) ? (src & (~XLEN_MASK)) + XLEN : src;
}

inline std::string RV64_RegName(RV64Reg reg) {
  std::string regname = "x" + std::to_string((int)reg);
  return regname;
}

inline RV64Reg RV64_RegOffset(RV64Reg reg, int offset) {
  auto xi = (int)reg + offset;
  return (RV64Reg)xi;
}

class MOperand;
class MInst;
class MBasicBlock;
class MFunction;
enum class MOpTag {
  Invalid,
  Virt,
  Imm,
  PreColor,
  Colored,
};
inline const char* enum_to_string(const MOpTag& tag) {
  const static char* _str_tab[] = { "Invalid", "Virt", "Imm", "PreColor", "Colored" };
  return _str_tab[(unsigned)tag];
}

class MOperand {
public:
  MOpTag op_type = MOpTag::Invalid;
  int value      = ~0;
  static MOperand VirtReg(int reg_no) {
    auto mop = MOperand({MOpTag::Virt, reg_no});
    return mop;
  }

  static MOperand Imm(int imm_val) {
    auto mop = MOperand({MOpTag::Imm, imm_val});
    return mop;
  }

  static MOperand PreClrReg(RV64Reg phy_reg) {
    auto mop = MOperand({MOpTag::PreColor, (int)phy_reg});
    return mop;
  }

  static MOperand AllocReg(RV64Reg phy_reg) {
    return MOperand({MOpTag::Colored, (int)phy_reg});
  }
  bool operator==(const MOperand &op2) const {
    return op_type == op2.op_type && value == op2.value;
  }
  bool operator!=(const MOperand &op2) const {
    return op_type != op2.op_type || value != op2.value;
  }
  bool operator<(const MOperand &op2) const {
    return op_type == op2.op_type ? value < op2.value : op_type < op2.op_type;
  }
  bool need_clr() const {
    return op_type == MOpTag::PreColor || op_type == MOpTag::Virt;
  }
  std::string to_string() {
    std::string ret = enum_to_string(op_type);
    ret += "(";
    switch (op_type) {
      case MOpTag::Imm:
      case MOpTag::Virt: ret += std::to_string(value); break;
      case MOpTag::PreColor: ret += enum_to_string((RV64Reg)value); break;
      case MOpTag::Colored: ret += enum_to_string((RV64Reg)value); break;
      case MOpTag::Invalid: assert(0);
    }
    ret += ")";
    return ret;
  }
};

enum class MInstTag {
  Add,
  Sub,
  Mul,
  Div,
  Mod,
  Lt,
  Le,
  Ge,
  Gt,
  Eq,
  Ne,
  And,
  Or,
  Lsh,  // sll
  Rsh,  // srl,sra
  Move, // actually a pseudo, mv = addi rt, rs, 0
  Branch,
  Jmp,
  Ret,
  Load,
  Store,
  Compare,
  Call,
  Globsym,
  Comment,
};

inline MInstTag inverse_cond(MInstTag src_tag) {
  switch (src_tag) {
    case MInstTag::Lt: return MInstTag::Ge;
    case MInstTag::Le: return MInstTag::Gt;
    case MInstTag::Ge: return MInstTag::Lt;
    case MInstTag::Gt: return MInstTag::Le;
    case MInstTag::Eq: return MInstTag::Ne;
    case MInstTag::Ne: return MInstTag::Eq;
    default: assert(0);
  }
}

class MInst {
public:
  MInstTag inst_tag;

  sptr(MBasicBlock) parent_bb;
  MInst(MInstTag tag, sptr(MBasicBlock) parent_bb) : inst_tag(tag), parent_bb(parent_bb) {}
  virtual ~MInst() = default;
};

class MBasicBlock {
public:
  sptr(BasicBlock) ir_bb;
  sptr(MFunction) parent_func;
  std::list<sptr(MInst)> inst_list;
  std::list<sptr(MBasicBlock)> pred_list;
  std::list<sptr(MBasicBlock)> succ_list;

  std::set<MOperand> use;
  std::set<MOperand> def;
  std::set<MOperand> livein;
  std::set<MOperand> liveout;
};

class MFunction {
public:
  std::list<sptr(MBasicBlock)> bb_list;
  sptr(Function) ir_func;

  std::list<sptr(MInst)> stack_arg_reloc;
  unsigned stack_size;

  unsigned virt_reg_cnt = 0;
};

class MCModule {
public:
  std::list<sptr(MFunction)> function_list;
  std::list<sptr(GlobalVar)> global_list;
  void IR2MC(const Module &ir_module);
};

class MInstBinary : public MInst {
public:
  MOperand dst;
  MOperand op1;
  MOperand op2;
  MInstBinary(MInstTag type, sptr(MBasicBlock) parent_bb) : MInst(type, parent_bb) {}

  static sptr(MInstBinary) New(MInstTag type, sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstBinary>(type, parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
};

class MInstJump : public MInst {
public:
  sptr(MBasicBlock) target;
  MInstJump(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Jmp, parent_bb) {}

  static sptr(MInstJump) New(sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstJump>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
};

class MInstBranch : public MInst {
public:
  sptr(MBasicBlock) target;
  MOperand op1;
  MOperand op2;
  MInstTag branch_tag;
  MInstBranch(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Branch, parent_bb) {}
  static sptr(MInstBranch) New(sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstBranch>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
};

class MInstLoad : public MInst {
public:
  MOperand dst;
  MOperand addr;
  MOperand offset;
  MInstLoad(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Load, parent_bb) {}

  static sptr(MInstLoad) New(sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstLoad>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
  static sptr(MInstLoad) New(sptr(MInst) rel_inst, bool insert_after = false) {
    auto parent_bb    = rel_inst->parent_bb;
    auto inst         = make_shared<MInstLoad>(parent_bb);
    auto rel_inst_itr = STD_FIND(parent_bb->inst_list, rel_inst);
    assert(rel_inst_itr != parent_bb->inst_list.end());
    if (insert_after) std::advance(rel_inst_itr, 1);
    parent_bb->inst_list.insert(rel_inst_itr, inst);
    return inst;
  }
};

class MInstStore : public MInst {
public:
  MOperand data;
  MOperand addr;
  MOperand offset;
  MInstStore(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Load, parent_bb) {}

  static sptr(MInstStore) New(sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstStore>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
  static sptr(MInstStore) New(sptr(MInst) rel_inst, bool insert_after = false) {
    auto parent_bb    = rel_inst->parent_bb;
    auto inst         = make_shared<MInstStore>(parent_bb);
    auto rel_inst_itr = STD_FIND(parent_bb->inst_list, rel_inst);
    assert(rel_inst_itr != parent_bb->inst_list.end());
    if (insert_after) std::advance(rel_inst_itr, 1);
    parent_bb->inst_list.insert(rel_inst_itr, inst);
    return inst;
  }
};

class MInstMove : public MInst {
public:
  MOperand dst;
  MOperand src;
  MInstMove(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Move, parent_bb) {}
  static sptr(MInstMove) New(sptr(MBasicBlock) parent_bb, bool insert_begin = false) {
    auto inst = make_shared<MInstMove>(parent_bb);
    if (insert_begin)
      parent_bb->inst_list.push_front(inst);
    else
      parent_bb->inst_list.push_back(inst);
    return inst;
  }
  static sptr(MInstMove) New(sptr(MInst) rel_inst, bool insert_after = false) {
    auto parent_bb    = rel_inst->parent_bb;
    auto inst         = make_shared<MInstMove>(parent_bb);
    auto rel_inst_itr = STD_FIND(parent_bb->inst_list, rel_inst);
    assert(rel_inst_itr != parent_bb->inst_list.end());
    if (insert_after) std::advance(rel_inst_itr, 1);
    parent_bb->inst_list.insert(rel_inst_itr, inst);
    return inst;
  }
};

class MInstSymbol : public MInst {
public:
  MOperand dst;
  sptr(GlobalVar) symbol;
  MInstSymbol(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Globsym, parent_bb) {}
  static sptr(MInstSymbol) New(sptr(MBasicBlock) parent_bb, bool insert_begin = false) {
    auto inst = make_shared<MInstSymbol>(parent_bb);
    if (insert_begin)
      parent_bb->inst_list.push_front(inst);
    else
      parent_bb->inst_list.push_back(inst);
    return inst;
  }
};

class MInstReturn : public MInst {
public:
  MInstReturn(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Ret, parent_bb) {}
  static sptr(MInstReturn) New(sptr(MBasicBlock) parent_bb, bool insert_begin = false) {
    auto inst = make_shared<MInstReturn>(parent_bb);
    if (insert_begin)
      parent_bb->inst_list.push_front(inst);
    else
      parent_bb->inst_list.push_back(inst);
    return inst;
  }
};

class MInstComment : public MInst {
public:
  std::string comment;
  MInstComment(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Comment, parent_bb) {}
  static sptr(MInstComment) New(sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstComment>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
};

class MInstCall : public MInst {
public:
  sptr(Function) ir_func;
  MInstCall(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Call, parent_bb) {}
  static sptr(MInstCall) New(sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstCall>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
};

void get_inst_defuse(sptr(MInst) inst, std::vector<MOperand> &def, std::vector<MOperand> &use);
void get_inst_defuse(sptr(MInst) inst, std::vector<MOperand *> &def, std::vector<MOperand *> &use);
void set_bb_def_use(sptr(MFunction) func);

} // namespace CompSysY