Use std::piecewise_constant_distribution instead of ad-hoc binary search.

[krasin]

Summary

Fix the issue with the most recently discovered unit receiving much less attention.

Note: this is the second attempt (prev: r258473). Now, libc++ build is fixed.

[kcc]

LGTM with nits

lib/Fuzzer/FuzzerInternal.h

204

clang-format, please

287

add a FIXME note saying that ideally we should be able to use GetRand()

[krasin]

Added a TODO.

[krasin]

clang-format

[krasin]

lib/Fuzzer/FuzzerInternal.h

204–205

Done.

287

Done. As discussed offline, the way to go is to expose the generator to the user supplied fuzzer, so that they could seed their internal PRNG, if needed.

lib/Fuzzer/FuzzerInternal.h

 #ifndef LLVM_FUZZER_INTERNAL_H
 #define LLVM_FUZZER_INTERNAL_H
 
 #include <cassert>
 #include <climits>
 #include <chrono>
 #include <cstddef>
 #include <cstdlib>
+#include <random>
 #include <string>
 #include <string.h>
 #include <vector>
 #include <unordered_set>
 
 #include "FuzzerInterface.h"
 
 namespace fuzzer {
struct FuzzingOptions {
     std::string ArtifactPrefix = "./";
     std::string ExactArtifactPath;
     bool SaveArtifacts = true;
     bool PrintNEW = true;  // Print a status line when new units are found;
     bool OutputCSV = false;
     bool PrintNewCovPcs = false;
   };
   Fuzzer(UserSuppliedFuzzer &USF, FuzzingOptions Options);
-  void AddToCorpus(const Unit &U) { Corpus.push_back(U); }
+  void AddToCorpus(const Unit &U) { Corpus.push_back(U); UpdateCorpusDistribution(); }
   size_t ChooseUnitIdxToMutate();
   const Unit &ChooseUnitToMutate() { return Corpus[ChooseUnitIdxToMutate()]; };
   void Loop();
   void Drill();
   void ShuffleAndMinimize();
   void InitializeTraceState();
   void AssignTaintLabels(uint8_t *Data, size_t Size);
   size_t CorpusSize() const { return Corpus.size(); }
   void ReadDir(const std::string &Path, long *Epoch) {
private:
   void MutateAndTestOne();
   void ReportNewCoverage(const Unit &U);
   bool RunOne(const Unit &U);
   void RunOneAndUpdateCorpus(Unit &U);
   void WriteToOutputCorpus(const Unit &U);
   void WriteUnitToFileWithPrefix(const Unit &U, const char *Prefix);
   void PrintStats(const char *Where, const char *End = "\n");
   void PrintStatusForNewUnit(const Unit &U);
+  // Updates the probability distribution for the units in the corpus.
+  // Must be called whenever the corpus or unit weights are changed.
+  void UpdateCorpusDistribution();
 
   void SyncCorpus();
 
   size_t RecordBlockCoverage();
   size_t RecordCallerCalleeCoverage();
   void PrepareCoverageBeforeRun();
   bool CheckCoverageAfterRun();
 
private:
   // For UseCounters
   std::vector<uint8_t> CounterBitmap;
   size_t TotalBits() {  // Slow. Call it only for printing stats.
     size_t Res = 0;
     for (auto x : CounterBitmap) Res += __builtin_popcount(x);
     return Res;
   }
 
+  // TODO(krasin): remove GetRand from UserSuppliedFuzzer,
+  // and fully rely on the generator and the seed.
+  // The user supplied fuzzer will have a way to access the
+  // generator for its own purposes (like seeding the custom
+  // PRNG).
+  std::mt19937 Generator;
+  std::piecewise_constant_distribution<double> CorpusDistribution;
   UserSuppliedFuzzer &USF;
   FuzzingOptions Options;
   system_clock::time_point ProcessStartTime = system_clock::now();
   system_clock::time_point LastExternalSync = system_clock::now();
   system_clock::time_point UnitStartTime;
   long TimeOfLongestUnitInSeconds = 0;
   long EpochOfLastReadOfOutputCorpus = 0;
   size_t LastRecordedBlockCoverage = 0;

lib/Fuzzer/FuzzerLoop.cpp

do {                                                                         \
     if (!fn)                                                                   \
       MissingWeakApiFunction(#fn);                                             \
   } while (false)
 
 // Only one Fuzzer per process.
 static Fuzzer *F;
 
 Fuzzer::Fuzzer(UserSuppliedFuzzer &USF, FuzzingOptions Options)
-    : USF(USF), Options(Options) {
+    : Generator(USF.GetRand().Rand()), USF(USF), Options(Options) {
   SetDeathCallback();
   InitializeTraceState();
   assert(!F);
   F = this;
 }
 
 void Fuzzer::SetDeathCallback() {
   CHECK_WEAK_API_FUNCTION(__sanitizer_set_death_callback);
void Fuzzer::RereadOutputCorpus() {
   if (Options.Verbosity >= 2)
     Printf("Reload: read %zd new units.\n",  AdditionalCorpus.size());
   for (auto &X : AdditionalCorpus) {
     if (X.size() > (size_t)Options.MaxLen)
       X.resize(Options.MaxLen);
     if (UnitHashesAddedToCorpus.insert(Hash(X)).second) {
       if (RunOne(X)) {
         Corpus.push_back(X);
+        UpdateCorpusDistribution();
         PrintStats("RELOAD");
       }
     }
   }
 }
 
 void Fuzzer::ShuffleAndMinimize() {
   bool PreferSmall = (Options.PreferSmallDuringInitialShuffle == 1 ||
for (size_t First = 0; First < 1; First++) {
       if (RunOne(U)) {
         NewCorpus.push_back(U);
         if (Options.Verbosity >= 2)
           Printf("NEW0: %zd L %zd\n", LastRecordedBlockCoverage, U.size());
       }
     }
   }
   Corpus = NewCorpus;
+  UpdateCorpusDistribution();
   for (auto &X : Corpus)
     UnitHashesAddedToCorpus.insert(Hash(X));
   PrintStats("INITED");
 }
 
 bool Fuzzer::RunOne(const Unit &U) {
   UnitStartTime = system_clock::now();
   TotalNumberOfRuns++;
if (Options.Verbosity) {
     Printf(" L: %zd ", U.size());
     USF.GetMD().PrintMutationSequence();
     Printf("\n");
   }
 }
 
 void Fuzzer::ReportNewCoverage(const Unit &U) {
   Corpus.push_back(U);
+  UpdateCorpusDistribution();
   UnitHashesAddedToCorpus.insert(Hash(U));
   USF.GetMD().RecordSuccessfulMutationSequence();
   PrintStatusForNewUnit(U);
   WriteToOutputCorpus(U);
   if (Options.ExitOnFirst)
     exit(0);
 }
 
if (i == 0)
       StartTraceRecording();
     RunOneAndUpdateCorpus(U);
     StopTraceRecording();
   }
 }
 
 // Returns an index of random unit from the corpus to mutate.
 // Hypothesis: units added to the corpus last are more likely to be interesting.
-// This function gives more wieght to the more recent units.
+// This function gives more weight to the more recent units.
 size_t Fuzzer::ChooseUnitIdxToMutate() {
-    size_t N = Corpus.size();
-    size_t Total = (N + 1) * N / 2;
-    size_t R = USF.GetRand()(Total);
-    size_t IdxBeg = 0, IdxEnd = N;
-    // Binary search.
-    while (IdxEnd - IdxBeg >= 2) {
-      size_t Idx = IdxBeg + (IdxEnd - IdxBeg) / 2;
-      if (R > (Idx + 1) * Idx / 2)
-        IdxBeg = Idx;
-      else
-        IdxEnd = Idx;
-    }
-    assert(IdxBeg < N);
-    return IdxBeg;
+    size_t Idx = static_cast<size_t>(CorpusDistribution(Generator));
+    assert(Idx < Corpus.size());
+    return Idx;
 }
 
 // Experimental search heuristic: drilling.
 // - Read, shuffle, execute and minimize the corpus.
 // - Choose one random unit.
 // - Reset the coverage.
 // - Start fuzzing as if the chosen unit was the only element of the corpus.
 // - When done, reset the coverage again.
 // - Merge the newly created corpus into the original one.
 void Fuzzer::Drill() {
   // The corpus is already read, shuffled, and minimized.
   assert(!Corpus.empty());
   Options.PrintNEW = false;  // Don't print NEW status lines when drilling.
 
   Unit U = ChooseUnitToMutate();
 
   CHECK_WEAK_API_FUNCTION(__sanitizer_reset_coverage);
   __sanitizer_reset_coverage();
 
   std::vector<Unit> SavedCorpus;
   SavedCorpus.swap(Corpus);
   Corpus.push_back(U);
+  UpdateCorpusDistribution();
   assert(Corpus.size() == 1);
   RunOne(U);
   PrintStats("DRILL ");
   std::string SavedOutputCorpusPath; // Don't write new units while drilling.
   SavedOutputCorpusPath.swap(Options.OutputCorpus);
   Loop();
 
   __sanitizer_reset_coverage();
void Fuzzer::SyncCorpus() {
   auto Now = system_clock::now();
   if (duration_cast<seconds>(Now - LastExternalSync).count() <
       Options.SyncTimeout)
     return;
   LastExternalSync = Now;
   ExecuteCommand(Options.SyncCommand + " " + Options.OutputCorpus);
 }
 
+void Fuzzer::UpdateCorpusDistribution() {
+  size_t N = Corpus.size();
+  std::vector<double> Intervals(N+1);
+  std::vector<double> Weights(N);
+  std::iota(Intervals.begin(), Intervals.end(), 0);
+  std::iota(Weights.begin(), Weights.end(), 1);
+  CorpusDistribution = std::piecewise_constant_distribution<double>(
+      Intervals.begin(), Intervals.end(), Weights.begin());
+}
+
 }  // namespace fuzzer

lib/Fuzzer/test/FuzzerUnittest.cpp

 #include "FuzzerInternal.h"
 #include "gtest/gtest.h"
 #include <set>
 
 using namespace fuzzer;
 
 // For now, have LLVMFuzzerTestOneInput just to make it link.
 // Later we may want to make unittests that actually call LLVMFuzzerTestOneInput.
-extern "C" void LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size) {
+extern "C" int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size) {
   abort();
 }
 
 TEST(Fuzzer, CrossOver) {
   FuzzerRandomLibc Rand(0);
   MutationDispatcher MD(Rand);
   Unit A({0, 1, 2}), B({5, 6, 7});
   Unit C;
TEST(FuzzerUtil, Base64) {
   EXPECT_EQ("YWI=", Base64({'a', 'b'}));
   EXPECT_EQ("eHk=", Base64({'x', 'y'}));
   EXPECT_EQ("YWJj", Base64({'a', 'b', 'c'}));
   EXPECT_EQ("eHl6", Base64({'x', 'y', 'z'}));
   EXPECT_EQ("YWJjeA==", Base64({'a', 'b', 'c', 'x'}));
   EXPECT_EQ("YWJjeHk=", Base64({'a', 'b', 'c', 'x', 'y'}));
   EXPECT_EQ("YWJjeHl6", Base64({'a', 'b', 'c', 'x', 'y', 'z'}));
 }
+
+TEST(Corpus, Distribution) {
+  FuzzerRandomLibc Rand(0);
+  SimpleUserSuppliedFuzzer USF(&Rand, LLVMFuzzerTestOneInput);
+  Fuzzer::FuzzingOptions Options;
+  Fuzzer Fuzz(USF, Options);
+  size_t N = 10;
+  size_t TriesPerUnit = 1<<20;
+  for (size_t i = 0; i < N; i++) {
+    Fuzz.AddToCorpus(Unit{ static_cast<uint8_t>(i) });
+  }
+  std::vector<size_t> Hist(N);
+  for (size_t i = 0; i < N * TriesPerUnit; i++) {
+    Hist[Fuzz.ChooseUnitIdxToMutate()]++;
+  }
+  for (size_t i = 0; i < N; i++) {
+    // A weak sanity check that every unit gets invoked.
+    EXPECT_GT(Hist[i], TriesPerUnit / N / 3);
+  }
+}