Pass RNG to functions that require randomness for cleaner interface.

PiperOrigin-RevId: 817646190

Author: xinhaoyuan

centipede/BUILD

@@ -603,6 +603,8 @@ cc_library(
         ":feature_set",
         ":runner_result",
         ":util",
+        "@abseil-cpp//absl/random",
+        "@abseil-cpp//absl/random:bit_gen_ref",
         "@abseil-cpp//absl/strings",
         "@com_google_fuzztest//common:defs",
         "@com_google_fuzztest//common:logging",

centipede/centipede.cc

@@ -840,9 +840,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,8 @@
 #include <utility>
 #include <vector>
 
+#include "absl/random/bit_gen_ref.h"
+#include "absl/random/random.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/substitute.h"
@@ -126,12 +128,12 @@ void Corpus::Add(const ByteArray& data, const FeatureVec& fv,
   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(absl::BitGenRef 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(absl::BitGenRef rng) const {
+  return records_[absl::Uniform<size_t>(rng, 0, records_.size())];
 }
 
 void Corpus::DumpStatsToFile(const FeatureSet &fs, std::string_view filepath,
@@ -210,18 +212,21 @@ void WeightedDistribution::RecomputeInternalState() {
 }
 
 __attribute__((noinline))  // to see it in profile.
-size_t
-WeightedDistribution::RandomIndex(size_t random) const {
+size_t WeightedDistribution::RandomIndex(absl::BitGenRef 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

@@ -23,6 +23,7 @@
 #include <utility>
 #include <vector>
 
+#include "absl/random/bit_gen_ref.h"
 #include "./centipede/binary_info.h"
 #include "./centipede/execution_metadata.h"
 #include "./centipede/feature.h"
@@ -49,7 +50,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(absl::BitGenRef rng) const;
   // Returns the number of weights.
   size_t size() const { return weights_.size(); }
   // Removes all weights.
@@ -131,9 +132,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(absl::BitGenRef rng) const;
   // Returns a random active corpus record using uniform distribution.
-  const CorpusRecord &UniformRandom(size_t random) const;
+  const CorpusRecord& UniformRandom(absl::BitGenRef 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

@@ -175,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();