[libFuzzer] more refactoring; don't compute sha1sum every time we mutate a unit from the corpus, use the stored one.

llvm-svn: 282115
diff --git a/llvm/lib/Fuzzer/FuzzerCorpus.h b/llvm/lib/Fuzzer/FuzzerCorpus.h
index 877a483..d42e7be 100644
--- a/llvm/lib/Fuzzer/FuzzerCorpus.h
+++ b/llvm/lib/Fuzzer/FuzzerCorpus.h
@@ -12,22 +12,26 @@
 #ifndef LLVM_FUZZER_CORPUS
 #define LLVM_FUZZER_CORPUS
 
+#include <random>
+
 #include "FuzzerDefs.h"
+#include "FuzzerRandom.h"
 
 namespace fuzzer {
 
 struct InputInfo {
   Unit U;  // The actual input data.
+  uint8_t Sha1[kSHA1NumBytes];  // Checksum.
 };
 
 class InputCorpus {
  public:
   InputCorpus() {
-    Corpus.reserve(1 << 14);  // Avoid too many resizes.
+    Inputs.reserve(1 << 14);  // Avoid too many resizes.
   }
-  size_t size() const { return Corpus.size(); }
-  bool empty() const { return Corpus.empty(); }
-  const Unit &operator[] (size_t Idx) const { return Corpus[Idx].U; }
+  size_t size() const { return Inputs.size(); }
+  bool empty() const { return Inputs.empty(); }
+  const Unit &operator[] (size_t Idx) const { return Inputs[Idx].U; }
   void Append(const std::vector<Unit> &V) {
     for (auto &U : V)
       push_back(U);
@@ -37,18 +41,47 @@
     if (!Hashes.insert(H).second) return;
     InputInfo II;
     II.U = U;
-    Corpus.push_back(II);
+    memcpy(II.Sha1, H.data(), kSHA1NumBytes);
+    Inputs.push_back(II);
+    UpdateCorpusDistribution();
   }
 
   typedef const std::vector<InputInfo>::const_iterator ConstIter;
-  ConstIter begin() const { return Corpus.begin(); }
-  ConstIter end() const { return Corpus.end(); }
+  ConstIter begin() const { return Inputs.begin(); }
+  ConstIter end() const { return Inputs.end(); }
 
   bool HasUnit(const Unit &U) { return Hashes.count(Hash(U)); }
+  const InputInfo &ChooseUnitToMutate(Random &Rand) {
+    return Inputs[ChooseUnitIdxToMutate(Rand)];
+  };
 
- private:
+  // 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 weight to the more recent units.
+  size_t ChooseUnitIdxToMutate(Random &Rand) {
+    size_t Idx =
+        static_cast<size_t>(CorpusDistribution(Rand.Get_mt19937()));
+    assert(Idx < Inputs.size());
+    return Idx;
+  }
+
+private:
+
+  // Updates the probability distribution for the units in the corpus.
+  // Must be called whenever the corpus or unit weights are changed.
+  void UpdateCorpusDistribution() {
+    size_t N = Inputs.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());
+  }
+  std::piecewise_constant_distribution<double> CorpusDistribution;
+
   std::unordered_set<std::string> Hashes;
-  std::vector<InputInfo> Corpus;
+  std::vector<InputInfo> Inputs;
 };
 
 }  // namespace fuzzer