Pass RNG to functions that require randomness for cleaner interface.

PiperOrigin-RevId: 817646190

Author: xinhaoyuan

centipede/BUILD

@@ -601,7 +601,9 @@ cc_library(
         ":execution_metadata",
         ":feature",
         ":feature_set",
+        ":runner_result",
         ":util",
+        "@abseil-cpp//absl/random",
         "@abseil-cpp//absl/strings",
         "@com_google_fuzztest//common:defs",
         "@com_google_fuzztest//common:logging",

centipede/centipede.cc

@@ -452,8 +452,8 @@ bool Centipede::RunBatch(
       batch_gained_new_coverage = true;
       FUZZTEST_CHECK_GT(fv.size(), 0UL);
       if (function_filter_passed) {
-        corpus_.Add(input_vec[i], fv, batch_result.results()[i].metadata(), fs_,
-                    coverage_frontier_);
+        corpus_.Add(input_vec[i], fv, batch_result.results()[i].metadata(),
+                    batch_result.results()[i].stats(), fs_, coverage_frontier_);
       }
       if (corpus_file != nullptr) {
         FUZZTEST_CHECK_OK(corpus_file->Write(input_vec[i]));
@@ -493,8 +493,9 @@ void Centipede::LoadShard(const Environment &load_env, size_t shard_index,
         FUZZTEST_VLOG(10) << "Adding input " << Hash(input)
                           << "; new features: " << num_new_features;
         fs_.MergeFeatures(input_features);
-        // TODO(kcc): cmp_args are currently not saved to disk and not reloaded.
-        corpus_.Add(input, input_features, {}, fs_, coverage_frontier_);
+        // TODO(xinhaoyuan): metadata and stats are currently not saved to disk
+        // and not reloaded.
+        corpus_.Add(input, input_features, {}, {}, fs_, coverage_frontier_);
         ++num_added_inputs;
       } else {
         FUZZTEST_VLOG(10) << "Skipping input: " << Hash(input);
@@ -759,7 +760,8 @@ void Centipede::LoadSeedInputs(BlobFileWriter *absl_nonnull corpus_file,
   // coverage and passed the filters.
   if (corpus_.NumTotal() == 0) {
     for (const auto &seed_input : seed_inputs)
-      corpus_.Add(seed_input, {}, {}, fs_, coverage_frontier_);
+      corpus_.Add(seed_input, /*fv=*/{}, /*metadata=*/{}, /*stats=*/{}, fs_,
+                  coverage_frontier_);
   }
 }
 
@@ -831,9 +833,9 @@ void Centipede::FuzzingLoop() {
     std::vector<MutationInputRef> mutation_inputs;
     mutation_inputs.reserve(env_.mutate_batch_size);
     for (size_t i = 0; i < env_.mutate_batch_size; i++) {
-      const auto &corpus_record = env_.use_corpus_weights
-                                      ? corpus_.WeightedRandom(rng_())
-                                      : corpus_.UniformRandom(rng_());
+      const auto& corpus_record = env_.use_corpus_weights
+                                      ? corpus_.WeightedRandom(rng_)
+                                      : corpus_.UniformRandom(rng_);
       mutation_inputs.push_back(
           MutationInputRef{corpus_record.data, &corpus_record.metadata});
     }

centipede/corpus.cc

@@ -22,6 +22,7 @@
 #include <utility>
 #include <vector>
 
+#include "absl/random/random.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/substitute.h"
@@ -30,6 +31,7 @@
 #include "./centipede/execution_metadata.h"
 #include "./centipede/feature.h"
 #include "./centipede/feature_set.h"
+#include "./centipede/runner_result.h"
 #include "./centipede/util.h"
 #include "./common/defs.h"
 #include "./common/logging.h"  // IWYU pragma: keep
@@ -113,23 +115,24 @@ size_t Corpus::Prune(const FeatureSet &fs,
   return subset_to_remove.size();
 }
 
-void Corpus::Add(const ByteArray &data, const FeatureVec &fv,
-                 const ExecutionMetadata &metadata, const FeatureSet &fs,
-                 const CoverageFrontier &coverage_frontier) {
+void Corpus::Add(const ByteArray& data, const FeatureVec& fv,
+                 const ExecutionMetadata& metadata,
+                 const ExecutionResult::Stats& stats, const FeatureSet& fs,
+                 const CoverageFrontier& coverage_frontier) {
   // TODO(kcc): use coverage_frontier.
   FUZZTEST_CHECK(!data.empty())
       << "Got request to add empty element to corpus: ignoring";
   FUZZTEST_CHECK_EQ(records_.size(), weighted_distribution_.size());
-  records_.push_back({data, fv, metadata});
+  records_.push_back({data, fv, metadata, stats});
   weighted_distribution_.AddWeight(ComputeWeight(fv, fs, coverage_frontier));
 }
 
-const CorpusRecord &Corpus::WeightedRandom(size_t random) const {
-  return records_[weighted_distribution_.RandomIndex(random)];
+const CorpusRecord& Corpus::WeightedRandom(Rng& rng) const {
+  return records_[weighted_distribution_.RandomIndex(rng)];
 }
 
-const CorpusRecord &Corpus::UniformRandom(size_t random) const {
-  return records_[random % records_.size()];
+const CorpusRecord& Corpus::UniformRandom(Rng& rng) const {
+  return records_[absl::Uniform<size_t>(rng, 0, records_.size())];
 }
 
 void Corpus::DumpStatsToFile(const FeatureSet &fs, std::string_view filepath,
@@ -208,18 +211,21 @@ void WeightedDistribution::RecomputeInternalState() {
 }
 
 __attribute__((noinline))  // to see it in profile.
-size_t
-WeightedDistribution::RandomIndex(size_t random) const {
+size_t WeightedDistribution::RandomIndex(Rng& rng) const {
   FUZZTEST_CHECK(!weights_.empty());
   FUZZTEST_CHECK(cumulative_weights_valid_);
-  uint64_t sum_of_all_weights = cumulative_weights_.back();
-  if (sum_of_all_weights == 0)
-    return random % size();  // can't do much else here.
-  random = random % sum_of_all_weights;
-  auto it = std::upper_bound(cumulative_weights_.begin(),
-                             cumulative_weights_.end(), random);
+  const uint64_t sum_of_all_weights = cumulative_weights_.back();
+  if (sum_of_all_weights == 0) {
+    // Can't do much else here.
+    return absl::Uniform<size_t>(rng, 0, size());
+  }
+  auto it =
+      std::upper_bound(cumulative_weights_.begin(), cumulative_weights_.end(),
+                       absl::Uniform<uint64_t>(rng, 0, sum_of_all_weights));
   FUZZTEST_CHECK(it != cumulative_weights_.end());
-  return it - cumulative_weights_.begin();
+  const size_t index = it - cumulative_weights_.begin();
+  FUZZTEST_CHECK(weights_[index] != 0);
+  return index;
 }
 
 uint64_t WeightedDistribution::PopBack() {

centipede/corpus.h

@@ -27,6 +27,7 @@
 #include "./centipede/execution_metadata.h"
 #include "./centipede/feature.h"
 #include "./centipede/feature_set.h"
+#include "./centipede/runner_result.h"
 #include "./centipede/util.h"
 #include "./common/defs.h"
 
@@ -48,7 +49,7 @@ class WeightedDistribution {
   // For proper randomness, `random` should come from a 64-bit RNG.
   // RandomIndex() must not be called after ChangeWeight() without first
   // calling RecomputeInternalState().
-  size_t RandomIndex(size_t random) const;
+  size_t RandomIndex(Rng& rng) const;
   // Returns the number of weights.
   size_t size() const { return weights_.size(); }
   // Removes all weights.
@@ -87,6 +88,7 @@ struct CorpusRecord {
   ByteArray data;
   FeatureVec features;
   ExecutionMetadata metadata;
+  ExecutionResult::Stats stats;
 };
 
 // Maintains the corpus of inputs.
@@ -105,9 +107,10 @@ class Corpus {
   // Adds a corpus element, consisting of 'data' (the input bytes, non-empty),
   // 'fv' (the features associated with this input), and execution `metadata`.
   // `fs` is used to compute weights of `fv`.
-  void Add(const ByteArray &data, const FeatureVec &fv,
-           const ExecutionMetadata &metadata, const FeatureSet &fs,
-           const CoverageFrontier &coverage_frontier);
+  void Add(const ByteArray& data, const FeatureVec& fv,
+           const ExecutionMetadata& metadata,
+           const ExecutionResult::Stats& stats, const FeatureSet& fs,
+           const CoverageFrontier& coverage_frontier);
   // Removes elements that contain only frequent features, according to 'fs'.
   // Also, randomly removes elements to reduce the size to <= `max_corpus_size`.
   // `max_corpus_size` should be positive.
@@ -128,9 +131,9 @@ class Corpus {
   std::pair<size_t, size_t> MaxAndAvgSize() const;
   // Returns a random active corpus record using weighted distribution.
   // See WeightedDistribution.
-  const CorpusRecord &WeightedRandom(size_t random) const;
+  const CorpusRecord& WeightedRandom(Rng& rng) const;
   // Returns a random active corpus record using uniform distribution.
-  const CorpusRecord &UniformRandom(size_t random) const;
+  const CorpusRecord& UniformRandom(Rng& rng) const;
   // Returns the element with index 'idx', where `idx` < NumActive().
   const ByteArray &Get(size_t idx) const { return records_[idx].data; }
   // Returns the execution metadata for the element `idx`, `idx` < NumActive().

centipede/corpus_test.cc

@@ -45,7 +45,8 @@ TEST(Corpus, GetCmpData) {
   ByteArray cmp_data{2, 0, 1, 2, 3};
   FeatureVec features1 = {10, 20, 30};
   fs.MergeFeatures(features1);
-  corpus.Add({1}, features1, /*metadata=*/{cmp_data}, fs, coverage_frontier);
+  corpus.Add({1}, features1, /*metadata=*/{cmp_data}, /*stats=*/{}, fs,
+             coverage_frontier);
   EXPECT_EQ(corpus.NumActive(), 1);
   EXPECT_EQ(corpus.GetMetadata(0).cmp_data, cmp_data);
 }
@@ -61,9 +62,9 @@ TEST(Corpus, PrintStats) {
   FeatureVec features1 = {10, 20, 30};
   FeatureVec features2 = {20, 40};
   fs.MergeFeatures(features1);
-  corpus.Add({1, 2, 3}, features1, {}, fs, coverage_frontier);
+  corpus.Add({1, 2, 3}, features1, {}, /*stats=*/{}, fs, coverage_frontier);
   fs.MergeFeatures(features2);
-  corpus.Add({4, 5}, features2, {}, fs, coverage_frontier);
+  corpus.Add({4, 5}, features2, {}, /*stats=*/{}, fs, coverage_frontier);
   const std::string stats_filepath = test_tmpdir / "corpus.txt";
   corpus.DumpStatsToFile(fs, stats_filepath, "Test corpus");
   std::string stats_file_contents;
@@ -93,7 +94,8 @@ TEST(Corpus, Prune) {
 
   auto Add = [&](const CorpusRecord& record) {
     fs.MergeFeatures(record.features);
-    corpus.Add(record.data, record.features, {}, fs, coverage_frontier);
+    corpus.Add(record.data, record.features, {}, /*stats=*/{}, fs,
+               coverage_frontier);
   };
 
   auto VerifyActiveInputs = [&](std::vector<ByteArray> expected_inputs) {
@@ -152,7 +154,8 @@ TEST(Corpus, PruneRegressionTest1) {
 
   auto Add = [&](const CorpusRecord& record) {
     fs.MergeFeatures(record.features);
-    corpus.Add(record.data, record.features, {}, fs, coverage_frontier);
+    corpus.Add(record.data, record.features, {}, /*stats=*/{}, fs,
+               coverage_frontier);
   };
 
   Add({{1}, {10, 20}});
@@ -172,20 +175,19 @@ TEST(WeightedDistribution, WeightedDistribution) {
     }
   };
 
+  Rng rng(12345);
   auto compute_freq = [&]() {
     freq.clear();
     freq.resize(wd.size());
-    // We use numbers in [0, kNumIter) instead of random numbers
-    // for simplicity.
     for (int i = 0; i < kNumIter; i++) {
-      freq[wd.RandomIndex(i)]++;
+      freq[wd.RandomIndex(rng)]++;
     }
   };
 
   set_weights({1, 1});
   compute_freq();
-  EXPECT_EQ(freq[0], kNumIter / 2);
-  EXPECT_EQ(freq[1], kNumIter / 2);
+  EXPECT_NEAR(freq[0], kNumIter / 2, 100);
+  EXPECT_NEAR(freq[1], kNumIter / 2, 100);
 
   set_weights({1, 2});
   compute_freq();
@@ -325,7 +327,7 @@ TEST(CoverageFrontier, Compute) {
 
   auto Add = [&](feature_t feature) {
     fs.MergeFeatures({feature});
-    corpus.Add({42}, {feature}, {}, fs, frontier);
+    corpus.Add({42}, {feature}, {}, /*stats=*/{}, fs, frontier);
   };
 
   // Add PC-based features.

centipede/util.h

@@ -121,12 +121,13 @@ template <typename T>
 void RemoveSubset(const std::vector<size_t> &subset_indices,
                   std::vector<T> &set) {
   size_t pos_to_write = 0;
-  for (size_t i = 0, n = set.size(); i < n; i++) {
-    // If subset_indices.size() is k, this loop's complexity is O(n*log(k)).
-    // We can do it in O(n+k) with a bit more code, but this loop is not
-    // expected to be hot. Besides, k would typically be small.
-    if (!std::binary_search(subset_indices.begin(), subset_indices.end(), i))
-      std::swap(set[pos_to_write++], set[i]);
+  for (size_t i = 0, n = set.size(), j = 0, m = subset_indices.size(); i < n;
+       i++) {
+    if (j < m && i == subset_indices[j]) {
+      ++j;
+      continue;
+    }
+    std::swap(set[pos_to_write++], set[i]);
   }
   set.resize(pos_to_write);
 }