#include "Instrument.h"

using namespace llvm;

namespace instrument {



/**
 * Implement your instrumentation for dynamic symbolic execution engine
 */
bool Instrument::runOnFunction(Function &F) {
    if(F.getName()!="main"){
        return false;
    }
    LLVMContext &context = F.getParent()->getContext();
    /// __DSE_Init__
    auto *FuncInitType = FunctionType::get(Type::getVoidTy(context),
                                           {},
                                           false);
    auto FuncInit = F.getParent()->getOrInsertFunction(DSE_Init, FuncInitType);

    /// __DSE_Alloca__
    auto *FuncAllocaType = FunctionType::get(Type::getVoidTy(context),
                                             {Type::getInt32Ty(context),Type::getInt32PtrTy(context)},
                                             false);
    auto FuncAlloca = F.getParent()->getOrInsertFunction(DSE_Alloca, FuncAllocaType);

    /// __DSE_Store__
    auto *FuncStoreType= FunctionType::get(Type::getVoidTy(context),
                                          {Type::getInt32PtrTy(context)},
                                          false);
    auto FuncStore = F.getParent()->getOrInsertFunction(DSE_Store, FuncStoreType);

    /// __DSE_Load__
    auto *FuncLoadType = FunctionType::get(Type::getVoidTy(context),
                                           {Type::getInt32Ty(context),Type::getInt32PtrTy(context)},
                                           false);
    auto FuncLoad = F.getParent()->getOrInsertFunction(DSE_Load, FuncLoadType);

    /// __DSE_ICmp__
    auto *FuncICmpType = FunctionType::get(Type::getVoidTy(context),
                                           {Type::getInt32Ty(context),Type::getInt32Ty(context)},
                                           false);
    auto FuncICmp = F.getParent()->getOrInsertFunction(DSE_ICmp, FuncICmpType);

    /// __DSE_Label__
    auto *FuncLabelType = FunctionType::get(Type::getVoidTy(context),
                                            {Type::getInt32Ty(context),
                                             Type::getInt32Ty(context),Type::getInt32Ty(context)},
                                            false);
    auto FuncLabel = F.getParent()->getOrInsertFunction(DSE_Label, FuncLabelType);

    /// __DSE_Loop__
    auto *FuncLoopType = FunctionType::get(Type::getVoidTy(context),
                                           {Type::getInt32Ty(context)},
                                           false);
    auto FuncLoop = F.getParent()->getOrInsertFunction(DSE_Loop, FuncLoopType);

    /// __DSE_Branch__
    auto *FuncBranchType = FunctionType::get(Type::getVoidTy(context),
                                             {Type::getInt32Ty(context),Type::getInt32Ty(context)},
                                             false);
    auto FuncBranch = F.getParent()->getOrInsertFunction(DSE_Branch, FuncBranchType);

    /// __DSE_BinOp__
    auto *FuncBinOpType = FunctionType::get(Type::getVoidTy(context),
                                            {Type::getInt32Ty(context),Type::getInt32Ty(context)},
                                            false);
    auto FuncBinOp = F.getParent()->getOrInsertFunction(DSE_BinOP, FuncBinOpType);

    /// __DSE_Const__
    auto *FuncConstType = FunctionType::get(Type::getVoidTy(context),
                                            {Type::getInt32Ty(context)},
                                            false);
    auto FuncConst = F.getParent()->getOrInsertFunction(DSE_Const,FuncConstType);

    /// __DSE_Register__
    auto *FuncRegisterType = FunctionType::get(Type::getVoidTy(context),
                                               {Type::getInt32Ty(context)},
                                               false);
    auto FuncRegister = F.getParent()->getOrInsertFunction(DSE_Register, FuncRegisterType);

     for(auto & B : F){
         if(B.getName().startswith("if") || B.getName().startswith("land.lhs") || B.getName().startswith("for")){
             CallInst::Create(FuncBranch,
                              {ConstantInt::get(Type::getInt32Ty(context), getRegisterID(&B)),
                               ConstantInt::get(Type::getInt32Ty(context), BranchID++)})
                              ->insertBefore(&B.front());
         }
         //errs() << "hello" <<  B.getName() << " " <<getRegisterID(&B) << '\n';
         for(auto &bb : B){
             auto bbName = bb.getName();
             //errs() << bbName << ":"<< getRegisterID(&bb) << '\n';
             if(bbName == "retval"){
                 CallInst::Create(FuncInit)
                                  ->insertBefore(&bb);
                 continue;
             }
             //errs() << bb.getOpcodeName() << ":"<< bb.getOpcode()<< '\n';
             switch (bb.getOpcode()) {
                 case RET:
                     break;
                 case BR:
                     if(dyn_cast<BranchInst>(&bb)->isConditional()){
                         //errs() << getRegisterID(bb.getOperand(2)) <<" "<< getBranchID(bb.getOperand(1))<<'\n';
                         CallInst::Create(FuncLabel,
                                          {ConstantInt::get(Type::getInt32Ty(context),
                                                            getRegisterID(bb.getOperand(0))),
                                                ConstantInt::get(Type::getInt32Ty(context),
                                                            getRegisterID(bb.getOperand(2))),
                                                ConstantInt::get(Type::getInt32Ty(context),1)})
                                                ->insertBefore(&bb);
                         CallInst::Create(FuncLabel,
                                          {ConstantInt::get(Type::getInt32Ty(context),
                                                            getRegisterID(bb.getOperand(0))),
                                                ConstantInt::get(Type::getInt32Ty(context),
                                                            getRegisterID(bb.getOperand(1))),
                                                ConstantInt::get(Type::getInt32Ty(context),0)})
                                                ->insertBefore(&bb);
                         CallInst::Create(FuncLoop,
                                          {ConstantInt::get(Type::getInt32Ty(context), 1)})
                                          ->insertBefore(&bb);
                     } else{
                         CallInst::Create(FuncLoop,
                                          {ConstantInt::get(Type::getInt32Ty(context), 0)})
                                          ->insertBefore(&bb);
                     }
                     break;
                 case ADD:
                 case SUB:
                 case MUL:
                 case SDIV:
                 case SREM:
                     CallInst::Create(FuncBinOp,
                                      {ConstantInt::get(Type::getInt32Ty(context), getRegisterID(&bb)),
                                       ConstantInt::get(Type::getInt32Ty(context),bb.getOpcode())})
                                       ->insertAfter(&bb);
                     for(unsigned i = 0; i < bb.getNumOperands(); i++){
                         if(dyn_cast<ConstantInt>(bb.getOperand(i))){
                             CallInst::Create(FuncConst,{ConstantInt::get(
                                     Type::getInt32Ty(context),
                                     dyn_cast<ConstantInt>(bb.getOperand(i))->getSExtValue())})
                                     ->insertAfter(&bb);
                         } else{
                             CallInst::Create(FuncRegister,{ConstantInt::get(Type::getInt32Ty(context),
                                                                             getRegisterID(bb.getOperand(i)))})
                                                                             ->insertAfter(&bb);
                         }
                     }
                     break;
                 case ALLOCA:
                     CallInst::Create(FuncAlloca,
                                      {ConstantInt::get(Type::getInt32Ty(context),getRegisterID(&bb)),
                                       &bb})
                                       ->insertAfter(&bb);

                     break;
                 case LOAD:
                     CallInst::Create(FuncLoad,
                                      {ConstantInt::get(Type::getInt32Ty(context),getRegisterID(&bb)),
                                       bb.getOperand(0)})
                                       ->insertAfter(&bb);
                     break;
                 case STORE:
                     if(dyn_cast<ConstantInt>(bb.getOperand(0))){
                         CallInst::Create(FuncStore,
                                          {bb.getOperand(1)})
                                 ->insertAfter(&bb);
                         CallInst::Create(FuncConst,
                                          {ConstantInt::get(
                                                  Type::getInt32Ty(context),
                                                  dyn_cast<ConstantInt>(bb.getOperand(0))->getSExtValue())})
                                                  ->insertAfter(&bb);
                     }
                     else{
                         CallInst::Create(FuncStore,
                                          {bb.getOperand(1)})
                                 ->insertAfter(&bb);
                         CallInst::Create(FuncRegister,
                                          {ConstantInt::get(Type::getInt32Ty(context),getRegisterID(bb.getOperand(0)))})
                                 ->insertAfter(&bb);

                     }
                     break;
                 case ICMP:
                     CallInst::Create(FuncICmp,
                                      {ConstantInt::get(Type::getInt32Ty(context),getRegisterID(&bb)),
                                       ConstantInt::get(Type::getInt32Ty(context),
                                                        dyn_cast<ICmpInst>(&bb)->getPredicate())})
                                                        ->insertAfter(&bb);
                     for(unsigned i = 0; i < bb.getNumOperands(); i++)
                     {
                         if(dyn_cast<ConstantInt>(bb.getOperand(i))){
                             CallInst::Create(FuncConst,
                                              {ConstantInt::get(
                                                      Type::getInt32Ty(context),
                                                      dyn_cast<ConstantInt>(bb.getOperand(i))->getSExtValue())})
                                     ->insertAfter(&bb);
                             //errs() << dyn_cast<ConstantInt>(bb.getOperand(i))->getSExtValue() <<'\n';
                         } else{
                             CallInst::Create(FuncRegister,
                                              {ConstantInt::get(Type::getInt32Ty(context),
                                                                getRegisterID(bb.getOperand(i)))})
                                     ->insertAfter(&bb);
                         }
                     }
                     break;
                 case CALL:
                 default:
                     break;
             }
         }
     }
    return false;
}



/// We initialize pass ID here. LLVM uses ID’s address to
/// identify a pass, so initialization value is not important.
char Instrument::ID = 1;
static RegisterPass<Instrument>
    /// Lastly, we register our class Instrument,giving it
    /// a command line argument “Instrument”, and a name
    /// “Instrumentation for Dynamic Symbolic Execution”.
    X("Instrument", "Instrumentation for Dynamic Symbolic Execution",
      false /* Only looks at CFG */,
      false /* Analysis Pass */);

} // namespace instrument
