No public description

PiperOrigin-RevId: 828592996

Author: xinhaoyuan

centipede/centipede.cc

@@ -133,6 +133,18 @@ Centipede::Centipede(const Environment &env, CentipedeCallbacks &user_callbacks,
   FUZZTEST_CHECK(env_.seed) << "env_.seed must not be zero";
   if (!env_.input_filter.empty() && env_.fork_server)
     input_filter_cmd_.StartForkServer(TemporaryLocalDirPath(), "input_filter");
+  if (env_.corpus_weight_method == Corpus::kWeightMethodNameForUniform) {
+    corpus_weight_method_ = Corpus::WeightMethod::Uniform;
+  } else if (env_.corpus_weight_method == Corpus::kWeightMethodNameForRecency) {
+    corpus_weight_method_ = Corpus::WeightMethod::Recency;
+  } else if (env_.corpus_weight_method == Corpus::kWeightMethodNameForRarity) {
+    corpus_weight_method_ = Corpus::WeightMethod::Rarity;
+  } else {
+    FUZZTEST_LOG(WARNING) << "Unknown corpus weight method "
+                          << env_.corpus_weight_method << " - fall back to "
+                          << Corpus::kWeightMethodNameForRarity;
+    corpus_weight_method_ = Corpus::WeightMethod::Rarity;
+  }
 }
 
 void Centipede::CorpusToFiles(const Environment &env, std::string_view dir) {
@@ -474,7 +486,8 @@ bool Centipede::RunBatch(
       }
     }
   }
-  corpus_.UpdateWeights(fs_, coverage_frontier_, env_.exec_time_weight_scaling);
+  corpus_.UpdateWeights(fs_, coverage_frontier_, corpus_weight_method_,
+                        env_.exec_time_weight_scaling);
   return batch_gained_new_coverage;
 }
 

centipede/centipede.h

@@ -185,6 +185,7 @@ class Centipede {
 
   FeatureSet fs_;
   Corpus corpus_;
+  Corpus::WeightMethod corpus_weight_method_;
   CoverageFrontier coverage_frontier_;
   size_t num_runs_ = 0;  // counts executed inputs
 

centipede/centipede_flags.inc

@@ -192,6 +192,9 @@ CENTIPEDE_FLAG(
     bool, use_corpus_weights, true,
     "If true, use weighted distribution when choosing the corpus element "
     "to mutate. This flag is mostly for Centipede developers.")
+CENTIPEDE_FLAG(std::string, corpus_weight_method, "rarity",
+               "The weight method to use on corpus. Available options are "
+               "`uniform`, `recency`, and `rarity` (fallback).")
 CENTIPEDE_FLAG(
     bool, exec_time_weight_scaling, true,
     "If true, scale the corpus weight by the execution time of each input.")

centipede/corpus.cc

@@ -45,15 +45,15 @@ namespace fuzztest::internal {
 //                                  Corpus
 //------------------------------------------------------------------------------
 
-// Returns the weight of `fv` computed using `fs` and `coverage_frontier`.
-static size_t ComputeWeight(const FeatureVec &fv, const FeatureSet &fs,
-                            const CoverageFrontier &coverage_frontier) {
-  size_t weight = fs.ComputeWeight(fv);
+// Returns the weight of `fv` computed using `coverage_frontier`.
+static size_t ComputeFrontierWeight(const FeatureVec& fv,
+                                    const CoverageFrontier& coverage_frontier) {
   // The following is checking for the cases where PCTable is not present. In
   // such cases, we cannot use any ControlFlow related features.
-  if (coverage_frontier.MaxPcIndex() == 0) return weight;
+  if (coverage_frontier.MaxPcIndex() == 0) return 1;
   size_t frontier_weights_sum = 0;
-  for (const auto feature : fv) {
+  for (size_t i = 0; i < fv.size(); ++i) {
+    const auto feature = fv[i];
     if (!feature_domains::kPCs.Contains(feature)) continue;
     const auto pc_index = ConvertPCFeatureToPcIndex(feature);
     // Avoid checking frontier for out-of-bounds indices.
@@ -63,7 +63,7 @@ static size_t ComputeWeight(const FeatureVec &fv, const FeatureSet &fs,
       frontier_weights_sum += coverage_frontier.FrontierWeight(pc_index);
     }
   }
-  return weight * (frontier_weights_sum + 1);  // Multiply by at least 1.
+  return frontier_weights_sum + 1;  // Multiply by at least 1.
 }
 
 std::pair<size_t, size_t> Corpus::MaxAndAvgSize() const {
@@ -79,14 +79,27 @@ std::pair<size_t, size_t> Corpus::MaxAndAvgSize() const {
 
 void Corpus::UpdateWeights(const FeatureSet& fs,
                            const CoverageFrontier& coverage_frontier,
-                           bool scale_by_exec_time) {
+                           WeightMethod method, bool scale_by_exec_time) {
   std::vector<double> weights;
   weights.resize(records_.size());
   for (size_t i = 0, n = records_.size(); i < n; ++i) {
     auto& record = records_[i];
     const size_t unseen = fs.PruneFeaturesAndCountUnseen(record.features);
     FUZZTEST_CHECK_EQ(unseen, 0);
-    weights[i] = fs.ComputeWeight(record.features);
+    switch (method) {
+      case WeightMethod::Uniform:
+        weights[i] = 1;
+        break;
+      case WeightMethod::Recency:
+        weights[i] = i;
+        break;
+      case WeightMethod::Rarity:
+        weights[i] = fs.ComputeRarityWeight(record.features);
+        break;
+      default:
+        FUZZTEST_LOG(FATAL) << "Unknown corpus weight method";
+    }
+    weights[i] *= ComputeFrontierWeight(record.features, coverage_frontier);
   }
   if (scale_by_exec_time) {
     double total_exec_time_usec = 0;
@@ -199,7 +212,8 @@ void Corpus::Add(const ByteArray& data, const FeatureVec& fv,
       << "Got request to add empty element to corpus: ignoring";
   FUZZTEST_CHECK_EQ(records_.size(), weighted_distribution_.size());
   records_.push_back({data, fv, metadata, stats});
-  weighted_distribution_.AddWeight(ComputeWeight(fv, fs, coverage_frontier));
+  // Will be updated by `UpdateWeights`.
+  weighted_distribution_.AddWeight(0);
 }
 
 const CorpusRecord& Corpus::WeightedRandom(absl::BitGenRef rng) const {

centipede/corpus.h

@@ -98,6 +98,16 @@ struct CorpusRecord {
 // Allows to prune (forget) inputs that become uninteresting.
 class Corpus {
  public:
+  enum class WeightMethod {
+    Uniform,
+    Recency,
+    Rarity,
+  };
+
+  static constexpr std::string_view kWeightMethodNameForUniform = "uniform";
+  static constexpr std::string_view kWeightMethodNameForRecency = "recency";
+  static constexpr std::string_view kWeightMethodNameForRarity = "rarity";
+
   Corpus() = default;
 
   Corpus(const Corpus &) = default;
@@ -120,12 +130,12 @@ class Corpus {
   // Returns the number of removed elements.
   size_t Prune(const FeatureSet &fs, const CoverageFrontier &coverage_frontier,
                size_t max_corpus_size, Rng &rng);
-  // Updates the corpus weights according to `fs` and `coverage_frontier`. If
-  // `scale_by_exec_time` is set, scales the weights by the corpus execution
-  // time relative to the average.
+  // Updates the corpus weights according to `fs` and `coverage_frontier` using
+  // the weight `method`. If `scale_by_exec_time` is set, scales the weights by
+  // the corpus execution time relative to the average.
   void UpdateWeights(const FeatureSet& fs,
                      const CoverageFrontier& coverage_frontier,
-                     bool scale_by_exec_time);
+                     WeightMethod method, bool scale_by_exec_time);
 
   // Accessors.
 

centipede/corpus_test.cc

@@ -114,7 +114,8 @@ TEST(Corpus, Prune) {
   Add({{2}, {30, 40}});
   Add({{3}, {40, 50}});
   Add({{4}, {10, 20}});
-  corpus.UpdateWeights(fs, coverage_frontier, /*scale_by_exec_time=*/false);
+  corpus.UpdateWeights(fs, coverage_frontier, Corpus::WeightMethod::Rarity,
+                       /*scale_by_exec_time=*/false);
 
   // Prune. Features 20 and 40 are frequent => input {0} will be removed.
   EXPECT_EQ(corpus.NumActive(), 5);
@@ -124,7 +125,8 @@ TEST(Corpus, Prune) {
   VerifyActiveInputs({{1}, {2}, {3}, {4}});
 
   Add({{5}, {30, 60}});
-  corpus.UpdateWeights(fs, coverage_frontier, /*scale_by_exec_time=*/false);
+  corpus.UpdateWeights(fs, coverage_frontier, Corpus::WeightMethod::Rarity,
+                       /*scale_by_exec_time=*/false);
 
   EXPECT_EQ(corpus.NumTotal(), 6);
   // Prune. Feature 30 is now frequent => inputs {1} and {2} will be removed.
@@ -181,14 +183,16 @@ TEST(Corpus, ScalesWeightsWithExecTime) {
   };
 
   // The weights should be equal without exec time scaling.
-  corpus.UpdateWeights(fs, coverage_frontier, /*scale_by_exec_time=*/false);
+  corpus.UpdateWeights(fs, coverage_frontier, Corpus::WeightMethod::Rarity,
+                       /*scale_by_exec_time=*/false);
   ComputeFreq();
   EXPECT_NEAR(freq[0], kNumIter / 3, 100);
   EXPECT_NEAR(freq[1], kNumIter / 3, 100);
   EXPECT_NEAR(freq[2], kNumIter / 3, 100);
 
   // The weights should favor {0} over {1} over {2} with exec time scaling.
-  corpus.UpdateWeights(fs, coverage_frontier, /*scale_by_exec_time=*/true);
+  corpus.UpdateWeights(fs, coverage_frontier, Corpus::WeightMethod::Rarity,
+                       /*scale_by_exec_time=*/true);
   ComputeFreq();
   EXPECT_GT(freq[0], freq[1] + 100);
   EXPECT_GT(freq[1], freq[2] + 100);

centipede/feature_set.cc

@@ -110,6 +110,7 @@ void FeatureSet::PruneDiscardedDomains(FeatureVec &features) const {
 }
 
 void FeatureSet::MergeFeatures(const FeatureVec& features) {
+  ++num_executions_;
   for (auto feature : features) {
     bool unseen = false;
     if (feature_domains::IsComparisonScoreFeature(feature)) {
@@ -129,18 +130,24 @@ void FeatureSet::MergeFeatures(const FeatureVec& features) {
       ++num_features_;
       ++features_per_domain_[feature_domains::Domain::FeatureToDomainId(
           feature)];
+      rare_features_.insert(feature);
     } else if (unseen) {
       freq = 0;
+      rare_features_.insert(feature);
     }
     if (freq < FrequencyThreshold(feature)) {
       ++freq;
+      if (freq >= FrequencyThreshold(feature)) {
+        rare_features_.erase(feature);
+      }
     }
   }
 }
 
 __attribute__((noinline))  // to see it in profile.
-uint64_t
-FeatureSet::ComputeWeight(const FeatureVec &features) const {
+double FeatureSet::ComputeRarityWeight(const FeatureVec& features) const {
+  // Use uint64_t to keep the previous behavior. Maybe we want to switch it to
+  // double later.
   uint64_t weight = 0;
   for (auto feature : features) {
     // The less frequent is the feature, the more valuable it is.
@@ -178,6 +185,7 @@ std::ostream &operator<<(std::ostream &out, const FeatureSet &fs) {
     out << " " << name << ": " << value;
   };
   out << "ft: " << fs.size();
+  out << " rare: " << fs.rare_features().size();
   LogIfNotZero(fs.CountFeatures(feature_domains::kPCs), "cov");
   LogIfNotZero(fs.CountFeatures(feature_domains::k8bitCounters), "cnt");
   LogIfNotZero(fs.CountFeatures(feature_domains::kDataFlow), "df");

centipede/feature_set.h

@@ -61,6 +61,12 @@ class FeatureSet {
   // How many different features are in the set.
   size_t size() const { return num_features_; }
 
+  size_t NumExecutions() const { return num_executions_; };
+
+  const std::unordered_set<feature_t>& rare_features() const {
+    return rare_features_;
+  }
+
   // Returns features that originate from CFG counters, converted to PCIndexVec.
   PCIndexVec ToCoveragePCs() const;
 
@@ -83,21 +89,22 @@ class FeatureSet {
   }
 
   // Returns the frequency associated with `feature`.
-  size_t Frequency(feature_t feature) const {
+  size_t Frequency(feature_t feature, bool skip_counter_match = false) const {
     if (feature_domains::IsComparisonScoreFeature(feature)) {
-      if ((feature & feature_domains::kCMPScoreBitmask) !=
-          cmp_scores_[feature_domains::CMPScoreFeatureIndex(feature)]) {
+      if (!skip_counter_match &&
+          (feature & feature_domains::kCMPScoreBitmask) !=
+              cmp_scores_[feature_domains::CMPScoreFeatureIndex(feature)]) {
         return 0;
       }
       feature &= ~feature_domains::kCMPScoreBitmask;
     }
     return frequencies_[feature];
   }
 
-  // Computes combined weight of `features`.
-  // The less frequent the feature is, the bigger its weight.
-  // The weight of a FeatureVec is a sum of individual feature weights.
-  uint64_t ComputeWeight(const FeatureVec &features) const;
+  // Computes combined weight of `features` based on the feature rarity that
+  // scales linearly. The less frequent the feature is, the bigger its
+  // weight. The weight of a FeatureVec is a sum of individual feature weights.
+  double ComputeRarityWeight(const FeatureVec& features) const;
 
   // Returns a debug string representing the state of *this.
   std::string DebugString() const;
@@ -135,6 +142,8 @@ class FeatureSet {
   // Unused parts of MmapNoReserveArray don't actually reserve memory.
   MmapNoReserveArray<kSize> frequencies_;
 
+  std::unordered_set<feature_t> rare_features_;
+
   static constexpr size_t kScoresSize =
       (feature_domains::kCMPScoreDomains.back().end() -
        feature_domains::kCMPScoreDomains.front().begin()) >>
@@ -146,6 +155,8 @@ class FeatureSet {
   // Counts all unique features added to this.
   size_t num_features_ = 0;
 
+  size_t num_executions_ = 0;
+
   // Counts features in each domain.
   size_t features_per_domain_[feature_domains::kNumDomains] = {};
 

centipede/feature_set_test.cc

@@ -27,8 +27,8 @@ namespace {
 TEST(FeatureSet, ComputeWeight) {
   FeatureSet feature_set(10, {});
 
-  auto W = [&](const FeatureVec &features) -> uint64_t {
-    return feature_set.ComputeWeight(features);
+  auto W = [&](const FeatureVec& features) -> uint64_t {
+    return feature_set.ComputeRarityWeight(features);
   };
 
   feature_set.MergeFeatures({1, 2, 3});
@@ -60,8 +60,8 @@ TEST(FeatureSet, ComputeWeightWithDifferentDomains) {
                              /* three features from domain #3 */ f3, f3 + 1,
                              f3 + 2});
 
-  auto weight = [&](const FeatureVec &features) -> uint64_t {
-    return feature_set.ComputeWeight(features);
+  auto weight = [&](const FeatureVec& features) -> uint64_t {
+    return feature_set.ComputeRarityWeight(features);
   };
 
   // Test that features from a less frequent domain have more weight.