slightly different weighting algo (#1719)

* better seed selection

* slightly different weighting calculation

* remove unnecessary memset

2 files changed, 65 insertions, 31 deletions

diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 831a0dbc..8fb7ecb1 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -1223,7 +1223,7 @@ double rand_next_percent(afl_state_t *afl);
 
 static inline u32 rand_below(afl_state_t *afl, u32 limit) {
 
-  if (limit <= 1) return 0;
+  if (unlikely(limit <= 1)) return 0;
 
   /* The boundary not being necessarily a power of 2,
      we need to ensure the result uniformity. */
@@ -1256,7 +1256,7 @@ static inline u32 rand_below(afl_state_t *afl, u32 limit) {
    expand havoc mode */
 static inline u32 rand_below_datalen(afl_state_t *afl, u32 limit) {
 
-  if (limit <= 1) return 0;
+  if (unlikely(limit <= 1)) return 0;
 
   switch (rand_below(afl, 3)) {
 
diff --git a/src/afl-fuzz-queue.c b/src/afl-fuzz-queue.c
index 8ad7cd97..b10bf749 100644
--- a/src/afl-fuzz-queue.c
+++ b/src/afl-fuzz-queue.c
@@ -49,11 +49,13 @@ inline u32 select_next_queue_entry(afl_state_t *afl) {
 
   u32    s = rand_below(afl, afl->queued_items);
   double p = rand_next_percent(afl);
+
   /*
   fprintf(stderr, "select: p=%f s=%u ... p < prob[s]=%f ? s=%u : alias[%u]=%u"
   " ==> %u\n", p, s, afl->alias_probability[s], s, s, afl->alias_table[s], p <
   afl->alias_probability[s] ? s : afl->alias_table[s]);
   */
+
   return (p < afl->alias_probability[s] ? s : afl->alias_table[s]);
 
 }
@@ -87,25 +89,28 @@ double compute_weight(afl_state_t *afl, struct queue_entry *q,
 
 void create_alias_table(afl_state_t *afl) {
 
-  u32    n = afl->queued_items, i = 0, a, g;
+  u32    n = afl->queued_items, i = 0, nSmall = 0, nLarge = n - 1;
   double sum = 0;
 
+  double *P = (double *)afl_realloc(AFL_BUF_PARAM(out), n * sizeof(double));
+  u32 *Small = (int *)afl_realloc(AFL_BUF_PARAM(out_scratch), n * sizeof(u32));
+  u32 *Large = (int *)afl_realloc(AFL_BUF_PARAM(in_scratch), n * sizeof(u32));
+
   afl->alias_table =
       (u32 *)afl_realloc((void **)&afl->alias_table, n * sizeof(u32));
   afl->alias_probability = (double *)afl_realloc(
       (void **)&afl->alias_probability, n * sizeof(double));
-  double *P = (double *)afl_realloc(AFL_BUF_PARAM(out), n * sizeof(double));
-  int    *S = (int *)afl_realloc(AFL_BUF_PARAM(out_scratch), n * sizeof(u32));
-  int    *L = (int *)afl_realloc(AFL_BUF_PARAM(in_scratch), n * sizeof(u32));
 
-  if (!P || !S || !L || !afl->alias_table || !afl->alias_probability) {
+  if (!P || !Small || !Large || !afl->alias_table || !afl->alias_probability) {
 
     FATAL("could not acquire memory for alias table");
 
   }
 
-  memset((void *)afl->alias_table, 0, n * sizeof(u32));
   memset((void *)afl->alias_probability, 0, n * sizeof(double));
+  memset((void *)afl->alias_table, 0, n * sizeof(u32));
+  memset((void *)Small, 0, n * sizeof(u32));
+  memset((void *)Large, 0, n * sizeof(u32));
 
   if (likely(afl->schedule < RARE)) {
 
@@ -166,7 +171,15 @@ void create_alias_table(afl_state_t *afl) {
     for (i = 0; i < n; i++) {
 
       // weight is always 0 for disabled entries
-      P[i] = (afl->queue_buf[i]->weight * n) / sum;
+      if (unlikely(afl->queue_buf[i]->disabled)) {
+
+        P[i] = 0;
+
+      } else {
+
+        P[i] = (afl->queue_buf[i]->weight * n) / sum;
+
+      }
 
     }
 
@@ -176,60 +189,81 @@ void create_alias_table(afl_state_t *afl) {
 
       struct queue_entry *q = afl->queue_buf[i];
 
-      if (likely(!q->disabled)) { q->perf_score = calculate_score(afl, q); }
+      if (likely(!q->disabled)) {
+
+        q->perf_score = calculate_score(afl, q);
+        sum += q->perf_score;
 
-      sum += q->perf_score;
+      }
 
     }
 
     for (i = 0; i < n; i++) {
 
       // perf_score is always 0 for disabled entries
-      P[i] = (afl->queue_buf[i]->perf_score * n) / sum;
+      if (unlikely(afl->queue_buf[i]->disabled)) {
+
+        P[i] = 0;
+
+      } else {
+
+        P[i] = (afl->queue_buf[i]->perf_score * n) / sum;
+
+      }
 
     }
 
   }
 
-  int nS = 0, nL = 0, s;
-  for (s = (s32)n - 1; s >= 0; --s) {
+  // Done collecting weightings in P, now create the arrays.
+
+  for (s32 j = (s32)(n - 1); j >= 0; j--) {
 
-    if (P[s] < 1) {
+    if (P[j] < 1) {
 
-      S[nS++] = s;
+      Small[nSmall++] = (u32)j;
 
     } else {
 
-      L[nL++] = s;
+      Large[nLarge--] = (u32)j;
 
     }
 
   }
 
-  while (nS && nL) {
+  while (nSmall && nLarge != n - 1) {
+
+    u32 small = Small[--nSmall];
+    u32 large = Large[++nLarge];
+
+    afl->alias_probability[small] = P[small];
+    afl->alias_table[small] = large;
 
-    a = S[--nS];
-    g = L[--nL];
-    afl->alias_probability[a] = P[a];
-    afl->alias_table[a] = g;
-    P[g] = P[g] + P[a] - 1;
-    if (P[g] < 1) {
+    P[large] = P[large] - (1 - P[small]);
 
-      S[nS++] = g;
+    if (P[large] < 1) {
+
+      Small[nSmall++] = large;
 
     } else {
 
-      L[nL++] = g;
+      Large[nLarge--] = large;
 
     }
 
   }
 
-  while (nL)
-    afl->alias_probability[L[--nL]] = 1;
+  while (nSmall) {
+
+    afl->alias_probability[Small[--nSmall]] = 1;
+
+  }
 
-  while (nS)
-    afl->alias_probability[S[--nS]] = 1;
+  while (nLarge != n - 1) {
+
+    afl->alias_probability[Large[++nLarge]] = 1;
+
+  }
 
   afl->reinit_table = 0;
 
@@ -264,7 +298,7 @@ void create_alias_table(afl_state_t *afl) {
   */
   /*
   fprintf(stderr, "  entry  alias  probability  perf_score   weight
-  filename\n"); for (u32 i = 0; i < n; ++i) fprintf(stderr, "  %5u  %5u  %11u
+  filename\n"); for (i = 0; i < n; ++i) fprintf(stderr, "  %5u  %5u  %11u
   %0.9f  %0.9f  %s\n", i, afl->alias_table[i], afl->alias_probability[i],
   afl->queue_buf[i]->perf_score, afl->queue_buf[i]->weight,
             afl->queue_buf[i]->fname);