1 files changed, 36 insertions, 23 deletions

diff --git a/lib/fuzzer/FuzzerCorpus.h b/lib/fuzzer/FuzzerCorpus.h
index 2da929835f45..8ad14656cffc 100644
--- a/lib/fuzzer/FuzzerCorpus.h
+++ b/lib/fuzzer/FuzzerCorpus.h
@@ -12,6 +12,7 @@
 #ifndef LLVM_FUZZER_CORPUS
 #define LLVM_FUZZER_CORPUS
 
+#include "FuzzerDataFlowTrace.h"
 #include "FuzzerDefs.h"
 #include "FuzzerIO.h"
 #include "FuzzerRandom.h"
@@ -35,8 +36,9 @@ struct InputInfo {
   size_t NumSuccessfullMutations = 0;
   bool MayDeleteFile = false;
   bool Reduced = false;
+  bool HasFocusFunction = false;
   Vector<uint32_t> UniqFeatureSet;
-  float FeatureFrequencyScore = 1.0;
+  Vector<uint8_t> DataFlowTraceForFocusFunction;
 };
 
 class InputCorpus {
@@ -45,7 +47,6 @@ class InputCorpus {
   InputCorpus(const std::string &OutputCorpus) : OutputCorpus(OutputCorpus) {
     memset(InputSizesPerFeature, 0, sizeof(InputSizesPerFeature));
     memset(SmallestElementPerFeature, 0, sizeof(SmallestElementPerFeature));
-    memset(FeatureFrequency, 0, sizeof(FeatureFrequency));
   }
   ~InputCorpus() {
     for (auto II : Inputs)
@@ -70,10 +71,24 @@ class InputCorpus {
         Res = std::max(Res, II->U.size());
     return Res;
   }
+
+  size_t NumInputsThatTouchFocusFunction() {
+    return std::count_if(Inputs.begin(), Inputs.end(), [](const InputInfo *II) {
+      return II->HasFocusFunction;
+    });
+  }
+
+  size_t NumInputsWithDataFlowTrace() {
+    return std::count_if(Inputs.begin(), Inputs.end(), [](const InputInfo *II) {
+      return !II->DataFlowTraceForFocusFunction.empty();
+    });
+  }
+
   bool empty() const { return Inputs.empty(); }
   const Unit &operator[] (size_t Idx) const { return Inputs[Idx]->U; }
   void AddToCorpus(const Unit &U, size_t NumFeatures, bool MayDeleteFile,
-                   const Vector<uint32_t> &FeatureSet) {
+                   bool HasFocusFunction, const Vector<uint32_t> &FeatureSet,
+                   const DataFlowTrace &DFT, const InputInfo *BaseII) {
     assert(!U.empty());
     if (FeatureDebug)
       Printf("ADD_TO_CORPUS %zd NF %zd\n", Inputs.size(), NumFeatures);
@@ -83,9 +98,19 @@ class InputCorpus {
     II.NumFeatures = NumFeatures;
     II.MayDeleteFile = MayDeleteFile;
     II.UniqFeatureSet = FeatureSet;
+    II.HasFocusFunction = HasFocusFunction;
     std::sort(II.UniqFeatureSet.begin(), II.UniqFeatureSet.end());
     ComputeSHA1(U.data(), U.size(), II.Sha1);
-    Hashes.insert(Sha1ToString(II.Sha1));
+    auto Sha1Str = Sha1ToString(II.Sha1);
+    Hashes.insert(Sha1Str);
+    if (HasFocusFunction)
+      if (auto V = DFT.Get(Sha1Str))
+        II.DataFlowTraceForFocusFunction = *V;
+    // This is a gross heuristic.
+    // Ideally, when we add an element to a corpus we need to know its DFT.
+    // But if we don't, we'll use the DFT of its base input.
+    if (II.DataFlowTraceForFocusFunction.empty() && BaseII)
+      II.DataFlowTraceForFocusFunction = BaseII->DataFlowTraceForFocusFunction;
     UpdateCorpusDistribution();
     PrintCorpus();
     // ValidateFeatureSet();
@@ -157,9 +182,9 @@ class InputCorpus {
   void PrintStats() {
     for (size_t i = 0; i < Inputs.size(); i++) {
       const auto &II = *Inputs[i];
-      Printf("  [%zd %s]\tsz: %zd\truns: %zd\tsucc: %zd\n", i,
+      Printf("  [% 3zd %s] sz: % 5zd runs: % 5zd succ: % 5zd focus: %d\n", i,
              Sha1ToString(II.Sha1).c_str(), II.U.size(),
-             II.NumExecutedMutations, II.NumSuccessfullMutations);
+             II.NumExecutedMutations, II.NumSuccessfullMutations, II.HasFocusFunction);
     }
   }
 
@@ -213,18 +238,10 @@ class InputCorpus {
     return false;
   }
 
-  void UpdateFeatureFrequency(size_t Idx) {
-    FeatureFrequency[Idx % kFeatureSetSize]++;
-  }
-  float GetFeatureFrequency(size_t Idx) const {
-    return FeatureFrequency[Idx % kFeatureSetSize];
-  }
-  void UpdateFeatureFrequencyScore(InputInfo *II) {
-    const float kMin = 0.01, kMax = 100.;
-    II->FeatureFrequencyScore = kMin;
-    for (auto Idx : II->UniqFeatureSet)
-      II->FeatureFrequencyScore += 1. / (GetFeatureFrequency(Idx) + 1.);
-    II->FeatureFrequencyScore = Min(II->FeatureFrequencyScore, kMax);
+  bool IsFeatureNew(size_t Idx, uint32_t NewSize, bool Shrink) {
+    assert(NewSize);
+    uint32_t OldSize = GetFeature(Idx % kFeatureSetSize);
+    return OldSize == 0 || (Shrink && OldSize > NewSize);
   }
 
   size_t NumFeatures() const { return NumAddedFeatures; }
@@ -264,14 +281,11 @@ private:
     std::iota(Intervals.begin(), Intervals.end(), 0);
     for (size_t i = 0; i < N; i++)
       Weights[i] = Inputs[i]->NumFeatures
-                       ? (i + 1) * Inputs[i]->FeatureFrequencyScore
+                       ? (i + 1) * (Inputs[i]->HasFocusFunction ? 1000 : 1)
                        : 0.;
     if (FeatureDebug) {
       for (size_t i = 0; i < N; i++)
         Printf("%zd ", Inputs[i]->NumFeatures);
-      Printf("NUM\n");
-      for (size_t i = 0; i < N; i++)
-        Printf("%f ", Inputs[i]->FeatureFrequencyScore);
       Printf("SCORE\n");
       for (size_t i = 0; i < N; i++)
         Printf("%f ", Weights[i]);
@@ -292,7 +306,6 @@ private:
   size_t NumUpdatedFeatures = 0;
   uint32_t InputSizesPerFeature[kFeatureSetSize];
   uint32_t SmallestElementPerFeature[kFeatureSetSize];
-  float FeatureFrequency[kFeatureSetSize];
 
   std::string OutputCorpus;
 };