4 files changed, 14 insertions, 15 deletions
diff --git a/src/afl-forkserver.c b/src/afl-forkserver.c
index 5cd000d7..6e1dfbba 100644
--- a/src/afl-forkserver.c
+++ b/src/afl-forkserver.c
@@ -701,10 +701,8 @@ void afl_fsrv_write_to_testcase(afl_forkserver_t *fsrv, u8 *buf, size_t len) {
 /* Execute target application, monitoring for timeouts. Return status
    information. The called program will update afl->fsrv->trace_bits. */
 
-fsrv_run_result_t afl_fsrv_run_target(
-    afl_forkserver_t *fsrv, u32 timeout,
-    void(classify_counts_func)(afl_forkserver_t *fsrv),
-    volatile u8 *stop_soon_p) {
+fsrv_run_result_t afl_fsrv_run_target(afl_forkserver_t *fsrv, u32 timeout,
+                                      volatile u8 *stop_soon_p) {
 
   s32 res;
   u32 exec_ms;
@@ -790,9 +788,6 @@ fsrv_run_result_t afl_fsrv_run_target(
      behave very normally and do not have to be treated as volatile. */
 
   MEM_BARRIER();
-  // u32 tb4 = *(u32 *)fsrv->trace_bits;
-
-  if (likely(classify_counts_func)) classify_counts_func(fsrv);
 
   /* Report outcome to caller. */
 
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index 3933acd8..594a9390 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -36,7 +36,10 @@
 fsrv_run_result_t run_target(afl_state_t *afl, afl_forkserver_t *fsrv,
                              u32 timeout) {
 
-  return afl_fsrv_run_target(fsrv, timeout, classify_counts, &afl->stop_soon);
+  fsrv_run_result_t res = afl_fsrv_run_target(fsrv, timeout, &afl->stop_soon);
+  // TODO: Don't classify for faults?
+  classify_counts(fsrv);
+  return res;
 
 }
 
diff --git a/src/afl-showmap.c b/src/afl-showmap.c
index 48436c34..97f377f3 100644
--- a/src/afl-showmap.c
+++ b/src/afl-showmap.c
@@ -223,13 +223,15 @@ void run_target_forkserver(afl_forkserver_t *fsrv, char **argv, u8 *mem,
 
   afl_fsrv_write_to_testcase(fsrv, mem, len);
 
-  if (afl_fsrv_run_target(fsrv, fsrv->exec_tmout, classify_counts,
-                          &stop_soon) == FSRV_RUN_ERROR) {
+  if (afl_fsrv_run_target(fsrv, fsrv->exec_tmout, &stop_soon) ==
+      FSRV_RUN_ERROR) {
 
     FATAL("Error running target");
 
   }
 
+  classify_counts(fsrv);
+
   if (stop_soon) {
 
     SAYF(cRST cLRD "\n+++ afl-showmap folder mode aborted by user +++\n" cRST);
diff --git a/src/afl-tmin.c b/src/afl-tmin.c
index cb53f56f..3330561b 100644
--- a/src/afl-tmin.c
+++ b/src/afl-tmin.c
@@ -121,8 +121,6 @@ static void apply_mask(u32 *mem, u32 *mask) {
 
 static void classify_counts(afl_forkserver_t *fsrv) {
 
-  if (hang_mode) return;                              /* We only want hangs */
-
   u8 *mem = fsrv->trace_bits;
   u32 i = MAP_SIZE;
 
@@ -146,8 +144,6 @@ static void classify_counts(afl_forkserver_t *fsrv) {
 
   }
 
-  apply_mask((u32 *)fsrv->trace_bits, (u32 *)mask_bitmap);
-
 }
 
 /* See if any bytes are set in the bitmap. */
@@ -224,7 +220,7 @@ static u8 run_target(afl_forkserver_t *fsrv, char **argv, u8 *mem, u32 len,
   afl_fsrv_write_to_testcase(fsrv, mem, len);
 
   fsrv_run_result_t ret =
-      afl_fsrv_run_target(fsrv, fsrv->exec_tmout, classify_counts, &stop_soon);
+      afl_fsrv_run_target(fsrv, fsrv->exec_tmout, &stop_soon);
 
   if (ret == FSRV_RUN_ERROR) FATAL("Couldn't run child");
 
@@ -250,6 +246,9 @@ static u8 run_target(afl_forkserver_t *fsrv, char **argv, u8 *mem, u32 len,
 
   }
 
+  classify_counts(fsrv);
+  apply_mask((u32 *)fsrv->trace_bits, (u32 *)mask_bitmap);
+
   if (ret == FSRV_RUN_TMOUT) {
 
     missed_hangs++;