#include "llvm/IR/Constants.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Module.h"
#include "llvm/Pass.h"
#include <map>

using namespace llvm;

namespace instrument {
#define RET         1
#define BR          2
#define ADD         13      ///    +
#define SUB         15      ///    -
#define MUL         17      ///    *
#define SDIV        20      ///    /
#define SREM        23      ///    %
#define ALLOCA      31
#define LOAD        32
#define STORE       33
#define ICMP        53
#define CALL        56

#define DSE_Init        "__DSE_Init__"
#define DSE_Alloca      "__DSE_Alloca__"
#define DSE_Store       "__DSE_Store__"
#define DSE_Load        "__DSE_Load__"
#define DSE_ICmp        "__DSE_ICmp__"
#define DSE_Label       "__DSE_Label__"
#define DSE_Loop        "__DSE_Loop__"
#define DSE_BinOP       "__DSE_BinOp__"
#define DSE_Branch      "__DSE_Branch__"
#define DSE_Const       "__DSE_Const__"
#define DSE_Register    "__DSE_Register__"


int RegisterID = 0;
std::map<Value *, int> RegisterMap;
int BranchID = 0;
std::map<Value *, int> BranchMap;

int getRegisterID(Value *I) {
  /// 找当前地址的ID，没有则分配一个
  if (RegisterMap.find(I) == RegisterMap.end()) {
    /// 为地址 *I 分配一个 ID
    RegisterMap[I] = RegisterID;
    return RegisterID++;
  } else {
    return RegisterMap[I];
  }
}

int getBranchID(Value *I) {
  if (BranchMap.find(I) == BranchMap.end()) {
    BranchMap[I] = BranchID;
    return BranchID++;
  } else {
    return BranchMap[I];
  }
}

/// 插桩，详细可以看官网文档
/// https://releases.llvm.org/14.0.0/docs/WritingAnLLVMPass.html
struct Instrument : public FunctionPass {
  static char ID;
  static const char *checkFunctionName;

  Instrument() : FunctionPass(ID) {}

  bool runOnFunction(Function &F) override;
};
} // namespace instrument