3 files changed, 16 insertions, 8 deletions
diff --git a/qemu_mode/patches/afl-qemu-cpu-inl.h b/qemu_mode/patches/afl-qemu-cpu-inl.h
index 03951fea..b769f62e 100644
--- a/qemu_mode/patches/afl-qemu-cpu-inl.h
+++ b/qemu_mode/patches/afl-qemu-cpu-inl.h
@@ -66,7 +66,7 @@ abi_ulong afl_entry_point, /* ELF entry point (_start) */
           afl_start_code,  /* .text start pointer      */
           afl_end_code;    /* .text end pointer        */
 
-u8 afl_enable_compcov;
+u8 afl_compcov_level;
 
 /* Set in the child process in forkserver mode: */
 
@@ -159,9 +159,14 @@ static void afl_setup(void) {
 
   }
   
+  /* Maintain for compatibility */
   if (getenv("AFL_QEMU_COMPCOV")) {
 
-    afl_enable_compcov = 1;
+    afl_compcov_level = 1;
+  }
+  if (getenv("AFL_COMPCOV_LEVEL")) {
+
+    afl_compcov_level = atoi(getenv("AFL_COMPCOV_LEVEL"));
   }
 
   /* pthread_atfork() seems somewhat broken in util/rcu.c, and I'm
diff --git a/qemu_mode/patches/afl-qemu-cpu-translate-inl.h b/qemu_mode/patches/afl-qemu-cpu-translate-inl.h
index 0ca89c98..4716c2ac 100644
--- a/qemu_mode/patches/afl-qemu-cpu-translate-inl.h
+++ b/qemu_mode/patches/afl-qemu-cpu-translate-inl.h
@@ -40,7 +40,7 @@
 extern unsigned char *afl_area_ptr;
 extern unsigned int afl_inst_rms;
 extern abi_ulong afl_start_code, afl_end_code;
-extern u8 afl_enable_compcov;
+extern u8 afl_compcov_level;
 
 void tcg_gen_afl_compcov_log_call(void *func, target_ulong cur_loc,
                                   TCGv_i64 arg1, TCGv_i64 arg2);
@@ -95,11 +95,14 @@ static void afl_compcov_log_64(target_ulong cur_loc, target_ulong arg1,
 
 
 static void afl_gen_compcov(target_ulong cur_loc, TCGv_i64 arg1, TCGv_i64 arg2,
-                            TCGMemOp ot) {
+                            TCGMemOp ot, int is_imm) {
 
   void *func;
   
-  if (!afl_enable_compcov || cur_loc > afl_end_code || cur_loc < afl_start_code)
+  if (!afl_compcov_level || cur_loc > afl_end_code || cur_loc < afl_start_code)
+    return;
+  
+  if (!is_imm && afl_compcov_level < 2)
     return;
 
   switch (ot) {
diff --git a/qemu_mode/patches/i386-translate.diff b/qemu_mode/patches/i386-translate.diff
index 0bc48828..239b2404 100644
--- a/qemu_mode/patches/i386-translate.diff
+++ b/qemu_mode/patches/i386-translate.diff
@@ -15,11 +15,11 @@ index 0dd5fbe4..b95d341e 100644
              tcg_gen_atomic_fetch_add_tl(s1->cc_srcT, s1->A0, s1->T0,
                                          s1->mem_index, ot | MO_LE);
              tcg_gen_sub_tl(s1->T0, s1->cc_srcT, s1->T1);
-+            afl_gen_compcov(s1->pc, s1->cc_srcT, s1->T1, ot);
++            afl_gen_compcov(s1->pc, s1->cc_srcT, s1->T1, ot, d == OR_EAX);
          } else {
              tcg_gen_mov_tl(s1->cc_srcT, s1->T0);
              tcg_gen_sub_tl(s1->T0, s1->T0, s1->T1);
-+            afl_gen_compcov(s1->pc, s1->T0, s1->T1, ot);
++            afl_gen_compcov(s1->pc, s1->T0, s1->T1, ot, d == OR_EAX);
              gen_op_st_rm_T0_A0(s1, ot, d);
          }
          gen_op_update2_cc(s1);
@@ -27,7 +27,7 @@ index 0dd5fbe4..b95d341e 100644
          tcg_gen_mov_tl(cpu_cc_src, s1->T1);
          tcg_gen_mov_tl(s1->cc_srcT, s1->T0);
          tcg_gen_sub_tl(cpu_cc_dst, s1->T0, s1->T1);
-+        afl_gen_compcov(s1->pc, s1->T0, s1->T1, ot);
++        afl_gen_compcov(s1->pc, s1->T0, s1->T1, ot, d == OR_EAX);
          set_cc_op(s1, CC_OP_SUBB + ot);
          break;
      }