1 files changed, 79 insertions, 30 deletions

diff --git a/src/afl-fuzz-queue.c b/src/afl-fuzz-queue.c
index fff8db03..48fd33ec 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]);
 
 }
@@ -74,7 +76,8 @@ double compute_weight(afl_state_t *afl, struct queue_entry *q,
   if (likely(afl->schedule < RARE)) { weight *= (avg_exec_us / q->exec_us); }
   weight *= (log(q->bitmap_size) / avg_bitmap_size);
   weight *= (1 + (q->tc_ref / avg_top_size));
-  if (unlikely(weight < 1.0)) { weight = 1.0; }
+
+  if (unlikely(weight < 0.1)) { weight = 0.1; }
   if (unlikely(q->favored)) { weight *= 5; }
   if (unlikely(!q->was_fuzzed)) { weight *= 2; }
 
@@ -86,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)) {
 
@@ -148,10 +154,32 @@ void create_alias_table(afl_state_t *afl) {
 
     }
 
+    if (unlikely(afl->schedule == MMOPT) && afl->queued_discovered) {
+
+      u32 cnt = afl->queued_discovered >= 5 ? 5 : afl->queued_discovered;
+
+      for (i = n - cnt; i < n; i++) {
+
+        struct queue_entry *q = afl->queue_buf[i];
+
+        if (likely(!q->disabled)) { q->weight *= 2.0; }
+
+      }
+
+    }
+
     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;
+
+      }
 
     }
 
@@ -161,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.
 
-    if (P[s] < 1) {
+  for (s32 j = (s32)(n - 1); j >= 0; j--) {
 
-      S[nS++] = s;
+    if (P[j] < 1) {
+
+      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 (nLarge != n - 1) {
 
-  while (nS)
-    afl->alias_probability[S[--nS]] = 1;
+    afl->alias_probability[Large[++nLarge]] = 1;
+
+  }
 
   afl->reinit_table = 0;
 
@@ -249,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);