Merge pull request #1137 from AFLplusplus/select

select instrumentation for coverage and cmplog

3 files changed, 603 insertions, 87 deletions

diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc
index 7a3d8c4d..ee8c317e 100644
--- a/instrumentation/SanitizerCoverageLTO.so.cc
+++ b/instrumentation/SanitizerCoverageLTO.so.cc
@@ -235,6 +235,8 @@ class ModuleSanitizerCoverage {
   uint32_t                         autodictionary = 1;
   uint32_t                         inst = 0;
   uint32_t                         afl_global_id = 0;
+  uint32_t                         unhandled = 0;
+  uint32_t                         select_cnt = 0;
   uint64_t                         map_addr = 0;
   const char *                     skip_nozero = NULL;
   const char *                     use_threadsafe_counters = nullptr;
@@ -1040,8 +1042,7 @@ bool ModuleSanitizerCoverage::instrumentModule(
           M, Int64Tyi, true, GlobalValue::ExternalLinkage, 0, "__afl_map_addr");
       ConstantInt *MapAddr = ConstantInt::get(Int64Tyi, map_addr);
       StoreInst *  StoreMapAddr = IRB.CreateStore(MapAddr, AFLMapAddrFixed);
-      StoreMapAddr->setMetadata(M.getMDKindID("nosanitize"),
-                                MDNode::get(Ctx, None));
+      ModuleSanitizerCoverage::SetNoSanitizeMetadata(StoreMapAddr);
 
     }
 
@@ -1056,8 +1057,7 @@ bool ModuleSanitizerCoverage::instrumentModule(
                              "__afl_final_loc");
       ConstantInt *const_loc = ConstantInt::get(Int32Tyi, write_loc);
       StoreInst *  StoreFinalLoc = IRB.CreateStore(const_loc, AFLFinalLoc);
-      StoreFinalLoc->setMetadata(M.getMDKindID("nosanitize"),
-                                 MDNode::get(Ctx, None));
+      ModuleSanitizerCoverage::SetNoSanitizeMetadata(StoreFinalLoc);
 
     }
 
@@ -1105,8 +1105,7 @@ bool ModuleSanitizerCoverage::instrumentModule(
                                0, "__afl_dictionary_len");
         ConstantInt *const_len = ConstantInt::get(Int32Tyi, offset);
         StoreInst *StoreDictLen = IRB.CreateStore(const_len, AFLDictionaryLen);
-        StoreDictLen->setMetadata(M.getMDKindID("nosanitize"),
-                                  MDNode::get(Ctx, None));
+        ModuleSanitizerCoverage::SetNoSanitizeMetadata(StoreDictLen);
 
         ArrayType *ArrayTy = ArrayType::get(IntegerType::get(Ctx, 8), offset);
         GlobalVariable *AFLInternalDictionary = new GlobalVariable(
@@ -1126,8 +1125,7 @@ bool ModuleSanitizerCoverage::instrumentModule(
         Value *AFLDictPtr =
             IRB.CreatePointerCast(AFLDictOff, PointerType::get(Int8Tyi, 0));
         StoreInst *StoreDict = IRB.CreateStore(AFLDictPtr, AFLDictionary);
-        StoreDict->setMetadata(M.getMDKindID("nosanitize"),
-                               MDNode::get(Ctx, None));
+        ModuleSanitizerCoverage::SetNoSanitizeMetadata(StoreDict);
 
       }
 
@@ -1150,9 +1148,9 @@ bool ModuleSanitizerCoverage::instrumentModule(
                getenv("AFL_USE_MSAN") ? ", MSAN" : "",
                getenv("AFL_USE_CFISAN") ? ", CFISAN" : "",
                getenv("AFL_USE_UBSAN") ? ", UBSAN" : "");
-      OKF("Instrumented %u locations with no collisions (on average %llu "
-          "collisions would be in afl-gcc/vanilla AFL) (%s mode).",
-          inst, calculateCollisions(inst), modeline);
+      OKF("Instrumented %u locations (%u selects) without collisions (%llu "
+          "collisions have been avoided) (%s mode).",
+          inst, select_cnt, calculateCollisions(inst), modeline);
 
     }
 
@@ -1274,6 +1272,7 @@ void ModuleSanitizerCoverage::instrumentFunction(
   const DominatorTree *    DT = DTCallback(F);
   const PostDominatorTree *PDT = PDTCallback(F);
   bool                     IsLeafFunc = true;
+  uint32_t                 skip_next = 0, local_selects = 0;
 
   for (auto &BB : F) {
 
@@ -1291,6 +1290,148 @@ void ModuleSanitizerCoverage::instrumentFunction(
 
         Value *val = ConstantInt::get(Int32Ty, ++afl_global_id);
         callInst->setOperand(1, val);
+        ++inst;
+
+      }
+
+      SelectInst *selectInst = nullptr;
+
+      /*
+            std::string errMsg;
+            raw_string_ostream os(errMsg);
+            IN.print(os);
+            fprintf(stderr, "X(%u): %s\n", skip_next, os.str().c_str());
+      */
+      if (!skip_next && (selectInst = dyn_cast<SelectInst>(&IN))) {
+
+        uint32_t    vector_cnt = 0;
+        Value *     condition = selectInst->getCondition();
+        Value *     result;
+        auto        t = condition->getType();
+        IRBuilder<> IRB(selectInst->getNextNode());
+
+        ++select_cnt;
+
+        if (t->getTypeID() == llvm::Type::IntegerTyID) {
+
+          Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id);
+          Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id);
+          result = IRB.CreateSelect(condition, val1, val2);
+          skip_next = 1;
+          inst += 2;
+
+        } else
+
+#if LLVM_VERSION_MAJOR > 13
+            if (t->getTypeID() == llvm::Type::FixedVectorTyID) {
+
+          FixedVectorType *tt = dyn_cast<FixedVectorType>(t);
+          if (tt) {
+
+            uint32_t elements = tt->getElementCount().getFixedValue();
+            vector_cnt = elements;
+            inst += vector_cnt * 2;
+            if (elements) {
+
+              FixedVectorType *GuardPtr1 =
+                  FixedVectorType::get(Int32Ty, elements);
+              FixedVectorType *GuardPtr2 =
+                  FixedVectorType::get(Int32Ty, elements);
+              Value *x, *y;
+
+              Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id);
+              Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id);
+              x = IRB.CreateInsertElement(GuardPtr1, val1, (uint64_t)0);
+              y = IRB.CreateInsertElement(GuardPtr2, val2, (uint64_t)0);
+
+              for (uint64_t i = 1; i < elements; i++) {
+
+                val1 = ConstantInt::get(Int32Ty, ++afl_global_id);
+                val2 = ConstantInt::get(Int32Ty, ++afl_global_id);
+                x = IRB.CreateInsertElement(GuardPtr1, val1, i);
+                y = IRB.CreateInsertElement(GuardPtr2, val2, i);
+
+              }
+
+              result = IRB.CreateSelect(condition, x, y);
+              skip_next = 1;
+
+            }
+
+          }
+
+        } else
+
+#endif
+        {
+
+          unhandled++;
+          continue;
+
+        }
+
+        local_selects++;
+        uint32_t vector_cur = 0;
+        /* Load SHM pointer */
+        LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr);
+        ModuleSanitizerCoverage::SetNoSanitizeMetadata(MapPtr);
+
+        while (1) {
+
+          /* Get CurLoc */
+          Value *MapPtrIdx = nullptr;
+
+          /* Load counter for CurLoc */
+          if (!vector_cnt) {
+
+            MapPtrIdx = IRB.CreateGEP(MapPtr, result);
+
+          } else {
+
+            auto element = IRB.CreateExtractElement(result, vector_cur++);
+            MapPtrIdx = IRB.CreateGEP(MapPtr, element);
+
+          }
+
+          if (use_threadsafe_counters) {
+
+            IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
+#if LLVM_VERSION_MAJOR >= 13
+                                llvm::MaybeAlign(1),
+#endif
+                                llvm::AtomicOrdering::Monotonic);
+
+          } else {
+
+            LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+            ModuleSanitizerCoverage::SetNoSanitizeMetadata(Counter);
+
+            /* Update bitmap */
+
+            Value *Incr = IRB.CreateAdd(Counter, One);
+
+            if (skip_nozero == NULL) {
+
+              auto cf = IRB.CreateICmpEQ(Incr, Zero);
+              auto carry = IRB.CreateZExt(cf, Int8Ty);
+              Incr = IRB.CreateAdd(Incr, carry);
+
+            }
+
+            auto nosan = IRB.CreateStore(Incr, MapPtrIdx);
+            ModuleSanitizerCoverage::SetNoSanitizeMetadata(nosan);
+
+          }
+
+          if (!vector_cnt || vector_cnt == vector_cur) { break; }
+
+        }
+
+        skip_next = 1;
+
+      } else {
+
+        skip_next = 0;
 
       }
 
@@ -1521,8 +1662,7 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
     } else {
 
       LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr);
-      MapPtr->setMetadata(Mo->getMDKindID("nosanitize"),
-                          MDNode::get(*Ct, None));
+      ModuleSanitizerCoverage::SetNoSanitizeMetadata(MapPtr);
       MapPtrIdx = IRB.CreateGEP(MapPtr, CurLoc);
 
     }
@@ -1539,8 +1679,7 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
     } else {
 
       LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
-      Counter->setMetadata(Mo->getMDKindID("nosanitize"),
-                           MDNode::get(*Ct, None));
+      ModuleSanitizerCoverage::SetNoSanitizeMetadata(Counter);
 
       Value *Incr = IRB.CreateAdd(Counter, One);
 
@@ -1552,8 +1691,8 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
 
       }
 
-      IRB.CreateStore(Incr, MapPtrIdx)
-          ->setMetadata(Mo->getMDKindID("nosanitize"), MDNode::get(*Ct, None));
+      auto nosan = IRB.CreateStore(Incr, MapPtrIdx);
+      ModuleSanitizerCoverage::SetNoSanitizeMetadata(nosan);
 
     }
 
diff --git a/instrumentation/SanitizerCoveragePCGUARD.so.cc b/instrumentation/SanitizerCoveragePCGUARD.so.cc
index 013492f9..be3f4f49 100644
--- a/instrumentation/SanitizerCoveragePCGUARD.so.cc
+++ b/instrumentation/SanitizerCoveragePCGUARD.so.cc
@@ -203,7 +203,7 @@ class ModuleSanitizerCoverage {
 
   SanitizerCoverageOptions Options;
 
-  uint32_t        instr = 0;
+  uint32_t        instr = 0, selects = 0, unhandled = 0;
   GlobalVariable *AFLMapPtr = NULL;
   ConstantInt *   One = NULL;
   ConstantInt *   Zero = NULL;
@@ -553,8 +553,9 @@ bool ModuleSanitizerCoverage::instrumentModule(
                getenv("AFL_USE_MSAN") ? ", MSAN" : "",
                getenv("AFL_USE_CFISAN") ? ", CFISAN" : "",
                getenv("AFL_USE_UBSAN") ? ", UBSAN" : "");
-      OKF("Instrumented %u locations with no collisions (%s mode).", instr,
-          modeline);
+      OKF("Instrumented %u locations with no collisions (%s mode) of which are "
+          "%u handled and %u unhandled selects.",
+          instr, modeline, selects, unhandled);
 
     }
 
@@ -833,9 +834,8 @@ bool ModuleSanitizerCoverage::InjectCoverage(Function &             F,
                                              ArrayRef<BasicBlock *> AllBlocks,
                                              bool IsLeafFunc) {
 
-  if (AllBlocks.empty()) return false;
+  uint32_t cnt_cov = 0, cnt_sel = 0, cnt_sel_inc = 0;
 
-  uint32_t special = 0;
   for (auto &BB : F) {
 
     for (auto &IN : BB) {
@@ -850,9 +850,37 @@ bool ModuleSanitizerCoverage::InjectCoverage(Function &             F,
         StringRef FuncName = Callee->getName();
         if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue;
 
-        uint32_t id = 1 + instr + (uint32_t)AllBlocks.size() + special++;
-        Value *  val = ConstantInt::get(Int32Ty, id);
-        callInst->setOperand(1, val);
+        cnt_cov++;
+
+      }
+
+      SelectInst *selectInst = nullptr;
+
+      if ((selectInst = dyn_cast<SelectInst>(&IN))) {
+
+        Value *c = selectInst->getCondition();
+        auto   t = c->getType();
+        if (t->getTypeID() == llvm::Type::IntegerTyID) {
+
+          cnt_sel++;
+          cnt_sel_inc += 2;
+
+        }
+
+#if LLVM__MAJOR > 11
+        else if (t->getTypeID() == llvm::Type::FixedVectorTyID) {
+
+          FixedVectorType *tt = dyn_cast<FixedVectorType>(t);
+          if (tt) {
+
+            cnt_sel++;
+            cnt_sel_inc += tt->getElementCount().getFixedValue();
+
+          }
+
+        }
+
+#endif
 
       }
 
@@ -860,11 +888,256 @@ bool ModuleSanitizerCoverage::InjectCoverage(Function &             F,
 
   }
 
-  CreateFunctionLocalArrays(F, AllBlocks, special);
-  for (size_t i = 0, N = AllBlocks.size(); i < N; i++)
-    InjectCoverageAtBlock(F, *AllBlocks[i], i, IsLeafFunc);
+  /* Create PCGUARD array */
+  CreateFunctionLocalArrays(F, AllBlocks, cnt_cov + cnt_sel_inc);
+  selects += cnt_sel;
+
+  uint32_t special = 0, local_selects = 0, skip_next = 0;
+
+  for (auto &BB : F) {
+
+    for (auto &IN : BB) {
+
+      CallInst *callInst = nullptr;
+
+      /*
+                                std::string errMsg;
+                                raw_string_ostream os(errMsg);
+                            IN.print(os);
+                            fprintf(stderr, "X: %s\n", os.str().c_str());
+      */
+      if ((callInst = dyn_cast<CallInst>(&IN))) {
+
+        Function *Callee = callInst->getCalledFunction();
+        if (!Callee) continue;
+        if (callInst->getCallingConv() != llvm::CallingConv::C) continue;
+        StringRef FuncName = Callee->getName();
+        if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue;
+
+        IRBuilder<> IRB(callInst);
+
+        Value *GuardPtr = IRB.CreateIntToPtr(
+            IRB.CreateAdd(
+                IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
+                ConstantInt::get(IntptrTy, (++special + AllBlocks.size()) * 4)),
+            Int32PtrTy);
+
+        LoadInst *Idx = IRB.CreateLoad(GuardPtr);
+        ModuleSanitizerCoverage::SetNoSanitizeMetadata(Idx);
+
+        callInst->setOperand(1, Idx);
+
+      }
+
+      SelectInst *selectInst = nullptr;
+
+      if (!skip_next && (selectInst = dyn_cast<SelectInst>(&IN))) {
+
+        uint32_t    vector_cnt = 0;
+        Value *     condition = selectInst->getCondition();
+        Value *     result;
+        auto        t = condition->getType();
+        IRBuilder<> IRB(selectInst->getNextNode());
+
+        if (t->getTypeID() == llvm::Type::IntegerTyID) {
+
+          auto GuardPtr1 = IRB.CreateIntToPtr(
+              IRB.CreateAdd(
+                  IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
+                  ConstantInt::get(
+                      IntptrTy,
+                      (cnt_cov + ++local_selects + AllBlocks.size()) * 4)),
+              Int32PtrTy);
+
+          auto GuardPtr2 = IRB.CreateIntToPtr(
+              IRB.CreateAdd(
+                  IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
+                  ConstantInt::get(
+                      IntptrTy,
+                      (cnt_cov + ++local_selects + AllBlocks.size()) * 4)),
+              Int32PtrTy);
+
+          result = IRB.CreateSelect(condition, GuardPtr1, GuardPtr2);
+
+        } else
+
+#if LLVM_VERSION_MAJOR > 13
+            if (t->getTypeID() == llvm::Type::FixedVectorTyID) {
+
+          FixedVectorType *tt = dyn_cast<FixedVectorType>(t);
+          if (tt) {
+
+            uint32_t elements = tt->getElementCount().getFixedValue();
+            vector_cnt = elements;
+            if (elements) {
+
+              FixedVectorType *GuardPtr1 =
+                  FixedVectorType::get(Int32PtrTy, elements);
+              FixedVectorType *GuardPtr2 =
+                  FixedVectorType::get(Int32PtrTy, elements);
+              Value *x, *y;
+
+              Value *val1 = IRB.CreateIntToPtr(
+                  IRB.CreateAdd(
+                      IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
+                      ConstantInt::get(
+                          IntptrTy,
+                          (cnt_cov + ++local_selects + AllBlocks.size()) * 4)),
+                  Int32PtrTy);
+              x = IRB.CreateInsertElement(GuardPtr1, val1, (uint64_t)0);
+
+              Value *val2 = IRB.CreateIntToPtr(
+                  IRB.CreateAdd(
+                      IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
+                      ConstantInt::get(
+                          IntptrTy,
+                          (cnt_cov + ++local_selects + AllBlocks.size()) * 4)),
+                  Int32PtrTy);
+              y = IRB.CreateInsertElement(GuardPtr2, val2, (uint64_t)0);
+
+              for (uint64_t i = 1; i < elements; i++) {
+
+                val1 = IRB.CreateIntToPtr(
+                    IRB.CreateAdd(
+                        IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
+                        ConstantInt::get(IntptrTy, (cnt_cov + ++local_selects +
+                                                    AllBlocks.size()) *
+                                                       4)),
+                    Int32PtrTy);
+                x = IRB.CreateInsertElement(x, val1, i);
+
+                val2 = IRB.CreateIntToPtr(
+                    IRB.CreateAdd(
+                        IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
+                        ConstantInt::get(IntptrTy, (cnt_cov + ++local_selects +
+                                                    AllBlocks.size()) *
+                                                       4)),
+                    Int32PtrTy);
+                y = IRB.CreateInsertElement(y, val2, i);
+
+              }
+
+              /*
+                          std::string errMsg;
+                          raw_string_ostream os(errMsg);
+                      x->print(os);
+                      fprintf(stderr, "X: %s\n", os.str().c_str());
+              */
+              result = IRB.CreateSelect(condition, x, y);
+
+            }
+
+          }
+
+        } else
+
+#endif
+        {
+
+          unhandled++;
+          continue;
+
+        }
+
+        local_selects++;
+        uint32_t vector_cur = 0;
 
-  instr += special;
+        /* Load SHM pointer */
+
+        LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr);
+        ModuleSanitizerCoverage::SetNoSanitizeMetadata(MapPtr);
+
+        /*
+            std::string errMsg;
+            raw_string_ostream os(errMsg);
+        result->print(os);
+        fprintf(stderr, "X: %s\n", os.str().c_str());
+        */
+
+        while (1) {
+
+          /* Get CurLoc */
+          LoadInst *CurLoc = nullptr;
+          Value *   MapPtrIdx = nullptr;
+
+          /* Load counter for CurLoc */
+          if (!vector_cnt) {
+
+            CurLoc = IRB.CreateLoad(result);
+            ModuleSanitizerCoverage::SetNoSanitizeMetadata(CurLoc);
+            MapPtrIdx = IRB.CreateGEP(MapPtr, CurLoc);
+
+          } else {
+
+            auto element = IRB.CreateExtractElement(result, vector_cur++);
+            auto elementptr = IRB.CreateIntToPtr(element, Int32PtrTy);
+            auto elementld = IRB.CreateLoad(elementptr);
+            ModuleSanitizerCoverage::SetNoSanitizeMetadata(elementld);
+            MapPtrIdx = IRB.CreateGEP(MapPtr, elementld);
+
+          }
+
+          if (use_threadsafe_counters) {
+
+            IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
+#if LLVM_VERSION_MAJOR >= 13
+                                llvm::MaybeAlign(1),
+#endif
+                                llvm::AtomicOrdering::Monotonic);
+
+          } else {
+
+            LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+            ModuleSanitizerCoverage::SetNoSanitizeMetadata(Counter);
+
+            /* Update bitmap */
+
+            Value *Incr = IRB.CreateAdd(Counter, One);
+
+            if (skip_nozero == NULL) {
+
+              auto cf = IRB.CreateICmpEQ(Incr, Zero);
+              auto carry = IRB.CreateZExt(cf, Int8Ty);
+              Incr = IRB.CreateAdd(Incr, carry);
+
+            }
+
+            StoreInst *StoreCtx = IRB.CreateStore(Incr, MapPtrIdx);
+            ModuleSanitizerCoverage::SetNoSanitizeMetadata(StoreCtx);
+
+          }
+
+          if (!vector_cnt) {
+
+            vector_cnt = 2;
+            break;
+
+          } else if (vector_cnt == vector_cur) {
+
+            break;
+
+          }
+
+        }
+
+        skip_next = 1;
+        instr += vector_cnt;
+
+      } else {
+
+        skip_next = 0;
+
+      }
+
+    }
+
+  }
+
+  if (AllBlocks.empty() && !special && !local_selects) return false;
+
+  if (!AllBlocks.empty())
+    for (size_t i = 0, N = AllBlocks.size(); i < N; i++)
+      InjectCoverageAtBlock(F, *AllBlocks[i], i, IsLeafFunc);
 
   return true;
 
@@ -1060,10 +1333,12 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
         Int32PtrTy);
 
     LoadInst *CurLoc = IRB.CreateLoad(GuardPtr);
+    ModuleSanitizerCoverage::SetNoSanitizeMetadata(CurLoc);
 
     /* Load SHM pointer */
 
     LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr);
+    ModuleSanitizerCoverage::SetNoSanitizeMetadata(MapPtr);
 
     /* Load counter for CurLoc */
 
@@ -1080,6 +1355,8 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
     } else {
 
       LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+      ModuleSanitizerCoverage::SetNoSanitizeMetadata(Counter);
+
       /* Update bitmap */
 
       Value *Incr = IRB.CreateAdd(Counter, One);
@@ -1092,7 +1369,8 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
 
       }
 
-      IRB.CreateStore(Incr, MapPtrIdx);
+      StoreInst *StoreCtx = IRB.CreateStore(Incr, MapPtrIdx);
+      ModuleSanitizerCoverage::SetNoSanitizeMetadata(StoreCtx);
 
     }
 
diff --git a/instrumentation/cmplog-instructions-pass.cc b/instrumentation/cmplog-instructions-pass.cc
index 0562c5b2..01a8a637 100644
--- a/instrumentation/cmplog-instructions-pass.cc
+++ b/instrumentation/cmplog-instructions-pass.cc
@@ -274,14 +274,15 @@ bool CmpLogInstructions::hookInstrs(Module &M) {
 
       Value *op0 = selectcmpInst->getOperand(0);
       Value *op1 = selectcmpInst->getOperand(1);
+      Value *op0_saved = op0, *op1_saved = op1;
+      auto   ty0 = op0->getType();
+      auto   ty1 = op1->getType();
 
-      IntegerType *        intTyOp0 = NULL;
-      IntegerType *        intTyOp1 = NULL;
-      unsigned             max_size = 0, cast_size = 0;
-      unsigned char        attr = 0;
-      std::vector<Value *> args;
-
-      CmpInst *cmpInst = dyn_cast<CmpInst>(selectcmpInst);
+      IntegerType *intTyOp0 = NULL;
+      IntegerType *intTyOp1 = NULL;
+      unsigned     max_size = 0, cast_size = 0;
+      unsigned     attr = 0, vector_cnt = 0;
+      CmpInst *    cmpInst = dyn_cast<CmpInst>(selectcmpInst);
 
       if (!cmpInst) { continue; }
 
@@ -327,7 +328,23 @@ bool CmpLogInstructions::hookInstrs(Module &M) {
 
       if (selectcmpInst->getOpcode() == Instruction::FCmp) {
 
-        auto ty0 = op0->getType();
+        if (ty0->isVectorTy()) {
+
+          VectorType *tt = dyn_cast<VectorType>(ty0);
+          if (!tt) {
+
+            fprintf(stderr, "Warning: cmplog cmp vector is not a vector!\n");
+            continue;
+
+          }
+
+#if LLVM_MAJOR > 11
+          vector_cnt = tt->getElementCount().getFixedValue();
+          ty0 = tt->getElementType();
+#endif
+
+        }
+
         if (ty0->isHalfTy()
 #if LLVM_VERSION_MAJOR >= 11
             || ty0->isBFloatTy()
@@ -342,13 +359,35 @@ bool CmpLogInstructions::hookInstrs(Module &M) {
           max_size = 80;
         else if (ty0->isFP128Ty() || ty0->isPPC_FP128Ty())
           max_size = 128;
+#if LLVM_MAJOR > 11
+        else if (ty0->getTypeID() != llvm::Type::PointerTyID && !be_quiet)
+          fprintf(stderr, "Warning: unsupported cmp type for cmplog: %u!\n",
+                  ty0->getTypeID());
+#endif
 
         attr += 8;
 
       } else {
 
-        intTyOp0 = dyn_cast<IntegerType>(op0->getType());
-        intTyOp1 = dyn_cast<IntegerType>(op1->getType());
+        if (ty0->isVectorTy()) {
+
+#if LLVM_MAJOR > 11
+          VectorType *tt = dyn_cast<VectorType>(ty0);
+          if (!tt) {
+
+            fprintf(stderr, "Warning: cmplog cmp vector is not a vector!\n");
+            continue;
+
+          }
+
+          vector_cnt = tt->getElementCount().getFixedValue();
+          ty1 = ty0 = tt->getElementType();
+#endif
+
+        }
+
+        intTyOp0 = dyn_cast<IntegerType>(ty0);
+        intTyOp1 = dyn_cast<IntegerType>(ty1);
 
         if (intTyOp0 && intTyOp1) {
 
@@ -356,11 +395,28 @@ bool CmpLogInstructions::hookInstrs(Module &M) {
                          ? intTyOp0->getBitWidth()
                          : intTyOp1->getBitWidth();
 
+        } else {
+
+#if LLVM_MAJOR > 11
+          if (ty0->getTypeID() != llvm::Type::PointerTyID && !be_quiet) {
+
+            fprintf(stderr, "Warning: unsupported cmp type for cmplog: %u\n",
+                    ty0->getTypeID());
+
+          }
+
+#endif
+
         }
 
       }
 
-      if (!max_size || max_size < 16) { continue; }
+      if (!max_size || max_size < 16) {
+
+        // fprintf(stderr, "too small\n");
+        continue;
+
+      }
 
       if (max_size % 8) { max_size = (((max_size / 8) + 1) * 8); }
 
@@ -393,67 +449,110 @@ bool CmpLogInstructions::hookInstrs(Module &M) {
 
       }
 
-      // errs() << "[CMPLOG] cmp  " << *cmpInst << "(in function " <<
-      // cmpInst->getFunction()->getName() << ")\n";
+      uint64_t cur = 0, last_val0 = 0, last_val1 = 0, cur_val;
 
-      // first bitcast to integer type of the same bitsize as the original
-      // type (this is a nop, if already integer)
-      Value *op0_i = IRB.CreateBitCast(
-          op0, IntegerType::get(C, op0->getType()->getPrimitiveSizeInBits()));
-      // then create a int cast, which does zext, trunc or bitcast. In our case
-      // usually zext to the next larger supported type (this is a nop if
-      // already the right type)
-      Value *V0 =
-          IRB.CreateIntCast(op0_i, IntegerType::get(C, cast_size), false);
-      args.push_back(V0);
-      Value *op1_i = IRB.CreateBitCast(
-          op1, IntegerType::get(C, op1->getType()->getPrimitiveSizeInBits()));
-      Value *V1 =
-          IRB.CreateIntCast(op1_i, IntegerType::get(C, cast_size), false);
-      args.push_back(V1);
+      while (1) {
 
-      // errs() << "[CMPLOG] casted parameters:\n0: " << *V0 << "\n1: " << *V1
-      // << "\n";
+        std::vector<Value *> args;
+        uint32_t             skip = 0;
 
-      ConstantInt *attribute = ConstantInt::get(Int8Ty, attr);
-      args.push_back(attribute);
+        if (vector_cnt) {
 
-      if (cast_size != max_size) {
+          op0 = IRB.CreateExtractElement(op0_saved, cur);
+          op1 = IRB.CreateExtractElement(op1_saved, cur);
+          ConstantInt *i0 = dyn_cast<ConstantInt>(op0);
+          ConstantInt *i1 = dyn_cast<ConstantInt>(op1);
+          if (i0 && i0->uge(0xffffffffffffffff) == false) {
 
-        ConstantInt *bitsize = ConstantInt::get(Int8Ty, (max_size / 8) - 1);
-        args.push_back(bitsize);
+            cur_val = i0->getZExtValue();
+            if (last_val0 && last_val0 == cur_val) { skip = 1; }
+            last_val0 = cur_val;
 
-      }
+          }
 
-      // fprintf(stderr, "_ExtInt(%u) castTo %u with attr %u didcast %u\n",
-      //         max_size, cast_size, attr);
+          if (i1 && i1->uge(0xffffffffffffffff) == false) {
 
-      switch (cast_size) {
+            cur_val = i1->getZExtValue();
+            if (last_val1 && last_val1 == cur_val) { skip = 1; }
+            last_val1 = cur_val;
 
-        case 8:
-          IRB.CreateCall(cmplogHookIns1, args);
-          break;
-        case 16:
-          IRB.CreateCall(cmplogHookIns2, args);
-          break;
-        case 32:
-          IRB.CreateCall(cmplogHookIns4, args);
-          break;
-        case 64:
-          IRB.CreateCall(cmplogHookIns8, args);
-          break;
-        case 128:
-          if (max_size == 128) {
+          }
+
+        }
+
+        if (!skip) {
+
+          // errs() << "[CMPLOG] cmp  " << *cmpInst << "(in function " <<
+          // cmpInst->getFunction()->getName() << ")\n";
+
+          // first bitcast to integer type of the same bitsize as the original
+          // type (this is a nop, if already integer)
+          Value *op0_i = IRB.CreateBitCast(
+              op0, IntegerType::get(C, ty0->getPrimitiveSizeInBits()));
+          // then create a int cast, which does zext, trunc or bitcast. In our
+          // case usually zext to the next larger supported type (this is a nop
+          // if already the right type)
+          Value *V0 =
+              IRB.CreateIntCast(op0_i, IntegerType::get(C, cast_size), false);
+          args.push_back(V0);
+          Value *op1_i = IRB.CreateBitCast(
+              op1, IntegerType::get(C, ty1->getPrimitiveSizeInBits()));
+          Value *V1 =
+              IRB.CreateIntCast(op1_i, IntegerType::get(C, cast_size), false);
+          args.push_back(V1);
 
-            IRB.CreateCall(cmplogHookIns16, args);
+          // errs() << "[CMPLOG] casted parameters:\n0: " << *V0 << "\n1: " <<
+          // *V1
+          // << "\n";
 
-          } else {
+          ConstantInt *attribute = ConstantInt::get(Int8Ty, attr);
+          args.push_back(attribute);
 
-            IRB.CreateCall(cmplogHookInsN, args);
+          if (cast_size != max_size) {
+
+            ConstantInt *bitsize = ConstantInt::get(Int8Ty, (max_size / 8) - 1);
+            args.push_back(bitsize);
 
           }
 
-          break;
+          // fprintf(stderr, "_ExtInt(%u) castTo %u with attr %u didcast %u\n",
+          //         max_size, cast_size, attr);
+
+          switch (cast_size) {
+
+            case 8:
+              IRB.CreateCall(cmplogHookIns1, args);
+              break;
+            case 16:
+              IRB.CreateCall(cmplogHookIns2, args);
+              break;
+            case 32:
+              IRB.CreateCall(cmplogHookIns4, args);
+              break;
+            case 64:
+              IRB.CreateCall(cmplogHookIns8, args);
+              break;
+            case 128:
+              if (max_size == 128) {
+
+                IRB.CreateCall(cmplogHookIns16, args);
+
+              } else {
+
+                IRB.CreateCall(cmplogHookInsN, args);
+
+              }
+
+              break;
+
+          }
+
+        }
+
+        /* else fprintf(stderr, "skipped\n"); */
+
+        ++cur;
+        if (cur >= vector_cnt) { break; }
 
       }