char *EncodeString(const char *Data, unsigned int lenght)
{
    char *NewData = (char *)malloc(lenght);
    for (unsigned int i = 0; i < lenght; i++) {
        NewData[i] = Data[i] + 1;
    }

    return NewData;
}

void decode(char *Data, unsigned int lenght)
{
    if (Data) {        
        for (unsigned int i = 0; i < lenght; i++) {
            Data[i] -= 1;
        }
    }
}

; Function Attrs: noinline nounwind optnone sspstrong uwtable
define dso_local void @decode(i8* %Data, i32 %lenght) #1 {
entry:
  %lenght.addr = alloca i32, align 4
  %Data.addr = alloca i8*, align 8
  %i = alloca i32, align 4
  store i32 %lenght, i32* %lenght.addr, align 4
  store i8* %Data, i8** %Data.addr, align 8
  %0 = load i8*, i8** %Data.addr, align 8
  %tobool = icmp ne i8* %0, null
  br i1 %tobool, label %if.then, label %if.end

if.then:                                          ; preds = %entry
  store i32 0, i32* %i, align 4
  br label %for.cond

for.cond:                                         ; preds = %for.inc, %if.then
  %1 = load i32, i32* %i, align 4
  %2 = load i32, i32* %lenght.addr, align 4
  %cmp = icmp ult i32 %1, %2
  br i1 %cmp, label %for.body, label %for.end

for.body:                                         ; preds = %for.cond
  %3 = load i8*, i8** %Data.addr, align 8
  %4 = load i32, i32* %i, align 4
  %idxprom = zext i32 %4 to i64
  %arrayidx = getelementptr inbounds i8, i8* %3, i64 %idxprom
  %5 = load i8, i8* %arrayidx, align 1
  %conv = sext i8 %5 to i32
  %sub = sub nsw i32 %conv, 1
  %conv1 = trunc i32 %sub to i8
  store i8 %conv1, i8* %arrayidx, align 1
  br label %for.inc

for.inc:                                          ; preds = %for.body
  %6 = load i32, i32* %i, align 4
  %inc = add i32 %6, 1
  store i32 %inc, i32* %i, align 4
  br label %for.cond

for.end:                                          ; preds = %for.cond
  br label %if.end

if.end:                                           ; preds = %for.end, %entry
  ret void
}

Function *createDecodeFunc(Module &M)
{
    auto &Ctx = M.getContext();

    // Add Decode function
    FunctionType *fn_type = FunctionType::get(Type::getVoidTy(Ctx), {Type::getInt8PtrTy(Ctx), Type::getInt32Ty(Ctx)}, false);
    FunctionCallee DecodeFuncCallee = M.getOrInsertFunction("decode", fn_type);
    Function *DecodeFunc = cast<Function>(DecodeFuncCallee.getCallee());
    DecodeFunc->setCallingConv(CallingConv::C);

    // Name DecodeFunc arguments
    Function::arg_iterator Args = DecodeFunc->arg_begin();
    Value *Data = Args++;
    Data->setName("Data");
    Value *lenght = Args++;
    lenght->setName("lenght");

    // Create blocks 
    BasicBlock *BEntry = BasicBlock::Create(Ctx, "entry", DecodeFunc);
    BasicBlock *BIfThen = BasicBlock::Create(Ctx, "if.then", DecodeFunc);
    BasicBlock *BIfEnd = BasicBlock::Create(Ctx, "if.end", DecodeFunc);
    BasicBlock *BForCond = BasicBlock::Create(Ctx, "for.cond", DecodeFunc);
    BasicBlock *BForBody = BasicBlock::Create(Ctx, "for.body", DecodeFunc);
    BasicBlock *BForInc = BasicBlock::Create(Ctx, "for.inc", DecodeFunc);
    BasicBlock *BForEnd = BasicBlock::Create(Ctx, "for.end", DecodeFunc);

    //// Entry block
    IRBuilder<> B1(BEntry);
    auto lenght_addr = B1.CreateAlloca(B1.getInt32Ty(), nullptr, "lenght.addr");
    auto Data_addr = B1.CreateAlloca(B1.getInt8PtrTy(), nullptr, "Data.addr");
    auto i = B1.CreateAlloca(B1.getInt32Ty(), nullptr, "i");

    B1.CreateStore(lenght, lenght_addr);
    B1.CreateStore(Data, Data_addr);

    auto var0 = B1.CreateLoad(Data_addr);
    auto tobool = B1.CreateICmpNE(var0, Constant::getNullValue(Type::getInt8PtrTy(Ctx)), "tobool");
    B1.CreateCondBr(tobool, BIfThen, BIfEnd);

    // if.then block
    IRBuilder<> B2(BIfThen);
    B2.CreateStore(B2.getInt32(0), i);
    B2.CreateBr(BForCond);

    // for.cond block
    IRBuilder<> B3(BForCond);
    auto var1 = B3.CreateLoad(i);
    auto var2 = B3.CreateLoad(lenght_addr);
    auto cmp = B3.CreateICmpULT(var1, var2, "cmp");
    B3.CreateCondBr(cmp, BForBody, BForEnd);

    // for.body block
    IRBuilder<> B4(BForBody);
    auto var3 = B4.CreateLoad(Data_addr);
    auto var4 = B4.CreateLoad(i);
    auto idxprom = B4.CreateZExt(var4, B4.getInt64Ty(), "idxprom");
    auto arrayidx = B4.CreateInBoundsGEP(B4.getInt8Ty(), var3, idxprom, "arrayidx");

    auto var5 = B4.CreateLoad(arrayidx);
    auto conv = B4.CreateSExt(var5, B4.getInt32Ty(), "conv");
    auto sub = B4.CreateSub(conv, B4.getInt32(1), "sub", false, true);
    auto conv1 = B4.CreateTrunc(sub, B4.getInt8Ty(), "conv1");
    B4.CreateStore(conv1, arrayidx);
    B4.CreateBr(BForInc);

    // for.inc block
    IRBuilder<> B5(BForInc);
    auto var6 = B5.CreateLoad(i);
    auto inc = B5.CreateAdd(var6, B5.getInt32(1), "inc");
    B5.CreateStore(inc, i);
    B5.CreateBr(BForCond);

    // for.end block
    IRBuilder<> B6(BForEnd);
    B6.CreateBr(BIfEnd);

    // End block
    IRBuilder<> B7(BIfEnd);
    B7.CreateRetVoid();

    return DecodeFunc;
}

Function *createDecodeStubFunc(Module &M, vector<GlobalString *> &GlobalStrings, Function *DecodeFunc)
{
    auto &Ctx = M.getContext();
    // Add DecodeStub function
    FunctionCallee DecodeStubCallee = M.getOrInsertFunction("decode_stub", Type::getVoidTy(Ctx));
    Function *DecodeStubFunc = cast<Function>(DecodeStubCallee.getCallee());
    DecodeStubFunc->setCallingConv(CallingConv::C);

    // Create entry block
    BasicBlock *BB = BasicBlock::Create(Ctx, "entry", DecodeStubFunc);
    IRBuilder<> Builder(BB);

    // Add calls to decode every encoded global
    for (GlobalString *GlobString : GlobalStrings) {
        if (GlobString->type == GlobString->SIMPLE_STRING_TYPE) {
            Constant *Zero = ConstantInt::get(Type::getInt32Ty(Ctx), 0);
            Value *Idx[] = {Zero, Zero};
            auto StrPtr = Builder.CreateInBoundsGEP(GlobString->Glob->getValueType(), GlobString->Glob, Idx);

            // Extract raw string
            auto CDA = dyn_cast<ConstantDataArray>(GlobString->Glob->getInitializer());
            StringRef StrVal = CDA->getAsString();
            size_t Size = StrVal.size();

            Builder.CreateCall(DecodeFunc, {StrPtr, ConstantInt::get(Builder.getInt32Ty(), Size)});
        }
        else if (GlobString->type == GlobString->STRUCT_STRING_TYPE) {
            auto StrPtr = Builder.CreateStructGEP(GlobString->Glob, GlobString->index);

            // Extract raw string
            auto CDA = dyn_cast<ConstantDataArray>(GlobString->Glob->getInitializer()->getOperand(GlobString->index));
            StringRef StrVal = CDA->getAsString();
            size_t Size = StrVal.size();

            Builder.CreateCall(DecodeFunc, {StrPtr, ConstantInt::get(Builder.getInt32Ty(), Size)});
        }
    }

    Builder.CreateRetVoid();

    // log IR info
    std::string str;
    raw_string_ostream OS(str);
    DecodeStubFunc->print(OS);
    errs() << str;

    return DecodeStubFunc;
}

void createDecodeStubBlock(Module &M, Function *DecodeStubFunc)
{
    Function *F = M.getFunction("main");
    auto &Ctx = F->getContext();
    BasicBlock &EntryBlock = F->getEntryBlock();

    auto Inst = EntryBlock.begin();
    // // Create new block
    // BasicBlock *NewBB = BasicBlock::Create(Ctx, "DecodeStub", EntryBlock.getParent(), &EntryBlock);
    // IRBuilder<> Builder(NewBB);
    // 
    // // Call stub func instruction
    // Builder.CreateCall(DecodeStubFunc);
    // // Jump to original entry block
    // Builder.CreateBr(&EntryBlock);

    IRBuilder<> Builder(&*Inst);
    Builder.CreateCall(DecodeStubFunc);
}