1 files changed, 13 insertions, 12 deletions
diff --git a/llvm_mode/afl-llvm-pass.so.cc b/llvm_mode/afl-llvm-pass.so.cc
index 9314c3d1..42a2f3af 100644
--- a/llvm_mode/afl-llvm-pass.so.cc
+++ b/llvm_mode/afl-llvm-pass.so.cc
@@ -84,7 +84,7 @@ class AFLCoverage : public ModulePass {
   uint32_t ngram_size = 0;
   uint32_t debug = 0;
   uint32_t map_size = MAP_SIZE;
-  char *   ctx_str = NULL;
+  char *   ctx_str = NULL, *skip_nozero = NULL;
 
 };
 
@@ -180,8 +180,9 @@ bool AFLCoverage::runOnModule(Module &M) {
 #if LLVM_VERSION_MAJOR < 9
   char *neverZero_counters_str = getenv("AFL_LLVM_NOT_ZERO");
 #endif
+  skip_nozero = getenv("AFL_LLVM_SKIP_NEVERZERO");
 
-  unsigned PrevLocSize;
+  unsigned PrevLocSize = 0;
 
   char *ngram_size_str = getenv("AFL_LLVM_NGRAM_SIZE");
   if (!ngram_size_str) ngram_size_str = getenv("AFL_NGRAM_SIZE");
@@ -215,9 +216,6 @@ bool AFLCoverage::runOnModule(Module &M) {
   if (ngram_size) PrevLocTy = VectorType::get(IntLocTy, PrevLocVecSize);
 #endif
 
-  if (ctx_str && ngram_size_str)
-    FATAL("you must decide between NGRAM and CTX instrumentation");
-
   /* Get globals for the SHM region and the previous location. Note that
      __afl_prev_loc is thread-local. */
 
@@ -436,8 +434,10 @@ bool AFLCoverage::runOnModule(Module &M) {
         PrevLocTrans = IRB.CreateXorReduce(PrevLoc);
       else
 #endif
-          if (ctx_str)
-        PrevLocTrans = IRB.CreateZExt(IRB.CreateXor(PrevLoc, PrevCtx), Int32Ty);
+        PrevLocTrans = PrevLoc;
+      if (ctx_str)
+        PrevLocTrans =
+            IRB.CreateZExt(IRB.CreateXor(PrevLocTrans, PrevCtx), Int32Ty);
       else
         PrevLocTrans = IRB.CreateZExt(PrevLoc, IRB.getInt32Ty());
 
@@ -451,7 +451,9 @@ bool AFLCoverage::runOnModule(Module &M) {
       if (ngram_size)
         MapPtrIdx = IRB.CreateGEP(
             MapPtr,
-            IRB.CreateZExt(IRB.CreateXor(PrevLocTrans, CurLoc), Int32Ty));
+            IRB.CreateZExt(
+                IRB.CreateXor(PrevLocTrans, IRB.CreateZExt(CurLoc, Int32Ty)),
+                Int32Ty));
       else
 #endif
         MapPtrIdx = IRB.CreateGEP(MapPtr, IRB.CreateXor(PrevLocTrans, CurLoc));
@@ -467,6 +469,9 @@ bool AFLCoverage::runOnModule(Module &M) {
       if (neverZero_counters_str !=
           NULL) {  // with llvm 9 we make this the default as the bug in llvm is
                    // then fixed
+#else
+      if (!skip_nozero) {
+
 #endif
         /* hexcoder: Realize a counter that skips zero during overflow.
          * Once this counter reaches its maximum value, it next increments to 1
@@ -482,12 +487,8 @@ bool AFLCoverage::runOnModule(Module &M) {
         auto carry = IRB.CreateZExt(cf, Int8Ty);
         Incr = IRB.CreateAdd(Incr, carry);
 
-#if LLVM_VERSION_MAJOR < 9
-
       }
 
-#endif
-
       IRB.CreateStore(Incr, MapPtrIdx)
           ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));