Record the origin for mutation for further tracking.

This is to enable further corpus management and scheduling methods based on the mutation stats.

For now, add the new `Mutant` type in :mutation_input. Will rename it to :mutation_data later.

PiperOrigin-RevId: 821630139

Author: xinhaoyuan

centipede/BUILD

@@ -439,6 +439,7 @@ cc_library(
         # used in centipede_runner.
         ":feature",
         ":execution_metadata",
+        ":mutation_input",
         ":shared_memory_blob_sequence",
         "@com_google_fuzztest//common:defs",
     ],
@@ -944,6 +945,7 @@ cc_library(
     hdrs = ["dispatcher.h"],
     deps = [
         ":execution_metadata",
+        ":mutation_input",
         ":runner_request",
         ":runner_result",
         ":shared_memory_blob_sequence",
@@ -1495,6 +1497,7 @@ cc_test(
     deps = [
         ":execution_metadata",
         ":feature",
+        ":mutation_input",
         ":runner_result",
         ":shared_memory_blob_sequence",
         "@com_google_fuzztest//common:defs",
@@ -1634,6 +1637,7 @@ cc_test(
         ":feature",
         ":feature_set",
         ":pc_info",
+        ":runner_result",
         ":util",
         "@com_google_fuzztest//common:defs",
         "@com_google_fuzztest//common:test_util",

centipede/byte_array_mutator.cc

@@ -321,27 +321,29 @@ void ByteArrayMutator::CrossOver(ByteArray &data, const ByteArray &other) {
 // TODO(kcc): add tests with different values of knobs.
 const KnobId knob_mutate_or_crossover = Knobs::NewId("mutate_or_crossover");
 
-std::vector<ByteArray> ByteArrayMutator::MutateMany(
-    const std::vector<MutationInputRef> &inputs, size_t num_mutants) {
+std::vector<Mutant> ByteArrayMutator::MutateMany(
+    const std::vector<MutationInputRef>& inputs, size_t num_mutants) {
   if (inputs.empty()) abort();
   // TODO(xinhaoyuan): Consider metadata in other inputs instead of always the
   // first one.
   SetMetadata(inputs[0].metadata != nullptr ? *inputs[0].metadata
                                             : ExecutionMetadata());
   size_t num_inputs = inputs.size();
-  std::vector<ByteArray> mutants;
+  std::vector<Mutant> mutants;
   mutants.reserve(num_mutants);
   for (size_t i = 0; i < num_mutants; ++i) {
-    auto mutant = inputs[rng_() % num_inputs].data;
-    if (mutant.size() <= max_len_ &&
+    Mutant mutant;
+    mutant.origin = rng_() % num_inputs;
+    mutant.data = inputs[mutant.origin].data;
+    if (mutant.data.size() <= max_len_ &&
         knobs_.GenerateBool(knob_mutate_or_crossover, rng_())) {
       // Do crossover only if the mutant is not over the max_len_.
       // Perform crossover with some other input. It may be the same input.
       const auto &other_input = inputs[rng_() % num_inputs].data;
-      CrossOver(mutant, other_input);
+      CrossOver(mutant.data, other_input);
     } else {
       // Perform mutation.
-      Mutate(mutant);
+      Mutate(mutant.data);
     }
     mutants.push_back(std::move(mutant));
   }

centipede/byte_array_mutator.h

@@ -108,8 +108,8 @@ class ByteArrayMutator {
   }
 
   // Takes non-empty `inputs` and produces `num_mutants` mutants.
-  std::vector<ByteArray> MutateMany(const std::vector<MutationInputRef> &inputs,
-                                    size_t num_mutants);
+  std::vector<Mutant> MutateMany(const std::vector<MutationInputRef>& inputs,
+                                 size_t num_mutants);
 
   using CrossOverFn = void (ByteArrayMutator::*)(ByteArray &,
                                                  const ByteArray &);

centipede/byte_array_mutator_test.cc

@@ -928,12 +928,12 @@ TEST(ByteArrayMutator, MutateManyWithAlignedInputs) {
       {0, 1, 2, 3, 4, 5, 6, 7},
       {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
   };
-  const std::vector<ByteArray> mutants =
+  const std::vector<Mutant> mutants =
       mutator.MutateMany(GetMutationInputRefsFromDataInputs(aligned_inputs),
                          kNumMutantsToGenerate);
   EXPECT_EQ(mutants.size(), kNumMutantsToGenerate);
-  for (const ByteArray &mutant : mutants) {
-    EXPECT_EQ(mutant.size() % kSizeAlignment, 0);
+  for (const Mutant& mutant : mutants) {
+    EXPECT_EQ(mutant.data.size() % kSizeAlignment, 0);
   }
 }
 
@@ -958,13 +958,13 @@ TEST(ByteArrayMutator, MutateManyWithUnalignedInputs) {
       {0, 1, 2, 3, 4, 5, 6, 7, 8, 9},
       {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
   };
-  const std::vector<ByteArray> mutants =
+  const std::vector<Mutant> mutants =
       mutator.MutateMany(GetMutationInputRefsFromDataInputs(unaligned_inputs),
                          kNumMutantsToGenerate);
   EXPECT_EQ(mutants.size(), kNumMutantsToGenerate);
-  for (const ByteArray &mutant : mutants) {
-    if (mutant.size() % kSizeAlignment != 0) {
-      EXPECT_LE(mutant.size(), 11);
+  for (const Mutant& mutant : mutants) {
+    if (mutant.data.size() % kSizeAlignment != 0) {
+      EXPECT_LE(mutant.data.size(), 11);
     }
   }
 }
@@ -982,12 +982,12 @@ TEST(ByteArrayMutator, MutateManyWithMaxLen) {
       {0, 1, 2},
       {0, 1, 2, 3},
   };
-  const std::vector<ByteArray> mutants = mutator.MutateMany(
+  const std::vector<Mutant> mutants = mutator.MutateMany(
       GetMutationInputRefsFromDataInputs(inputs), kNumMutantsToGenerate);
   EXPECT_EQ(mutants.size(), kNumMutantsToGenerate);
 
-  for (const ByteArray &mutant : mutants) {
-    EXPECT_LE(mutant.size(), kMaxLen);
+  for (const Mutant& mutant : mutants) {
+    EXPECT_LE(mutant.data.size(), kMaxLen);
   }
 }
 
@@ -1001,16 +1001,16 @@ TEST(ByteArrayMutator, MutateManyWithMaxLenWithStartingLargeInput) {
   const std::vector<ByteArray> large_input = {
       {0, 1, 2, 3, 4, 5, 6, 7}, {0}, {0, 1}, {0, 1, 2}, {0, 1, 2, 3},
   };
-  const std::vector<ByteArray> mutants = mutator.MutateMany(
+  const std::vector<Mutant> mutants = mutator.MutateMany(
       GetMutationInputRefsFromDataInputs(large_input), kNumMutantsToGenerate);
   EXPECT_EQ(mutants.size(), kNumMutantsToGenerate);
 
-  for (const ByteArray &mutant : mutants) {
-    if (mutant.size() > kMaxLen) {
+  for (const Mutant& mutant : mutants) {
+    if (mutant.data.size() > kMaxLen) {
       // The only mutant larger than max length should be the same large input
       // that mutation originally started with. All other mutants should be
       // within the maximum length specified.
-      EXPECT_EQ(mutant, large_input[0]);
+      EXPECT_EQ(mutant.data, large_input[0]);
     }
   }
 }

centipede/centipede.cc

@@ -416,13 +416,18 @@ size_t Centipede::AddPcPairFeatures(FeatureVec &fv) {
 }
 
 bool Centipede::RunBatch(
-    const std::vector<ByteArray> &input_vec,
-    BlobFileWriter *absl_nullable corpus_file,
-    BlobFileWriter *absl_nullable features_file,
-    BlobFileWriter *absl_nullable unconditional_features_file) {
+    const std::vector<ByteArray>& input_vec,
+    const std::vector<size_t>& mutant_origins,
+    BlobFileWriter* absl_nullable corpus_file,
+    BlobFileWriter* absl_nullable features_file,
+    BlobFileWriter* absl_nullable unconditional_features_file) {
   BatchResult batch_result;
   bool success = ExecuteAndReportCrash(env_.binary, input_vec, batch_result);
   FUZZTEST_CHECK_EQ(input_vec.size(), batch_result.results().size());
+  FUZZTEST_CHECK(mutant_origins.empty() ||
+                 mutant_origins.size() >= input_vec.size())
+      << "Got " << mutant_origins.size() << " with " << input_vec.size()
+      << " input";
 
   for (const auto &extra_binary : env_.extra_binaries) {
     if (ShouldStop()) break;
@@ -474,6 +479,7 @@ bool Centipede::RunBatch(
       }
     }
   }
+  corpus_.UpdateWeights(fs_, coverage_frontier_, env_.exec_time_weight_scaling);
   return batch_gained_new_coverage;
 }
 
@@ -563,7 +569,7 @@ void Centipede::Rerun(std::vector<ByteArray> &to_rerun) {
     size_t batch_size = std::min(to_rerun.size(), env_.batch_size);
     std::vector<ByteArray> batch(to_rerun.end() - batch_size, to_rerun.end());
     to_rerun.resize(to_rerun.size() - batch_size);
-    if (RunBatch(batch, nullptr, nullptr, features_file.get())) {
+    if (RunBatch(batch, {}, nullptr, nullptr, features_file.get())) {
       UpdateAndMaybeLogStats("rerun-old", 1);
     }
   }
@@ -757,7 +763,7 @@ void Centipede::LoadSeedInputs(BlobFileWriter *absl_nonnull corpus_file,
     seed_inputs.push_back({0});
   }
 
-  RunBatch(seed_inputs, corpus_file, features_file,
+  RunBatch(seed_inputs, {}, corpus_file, features_file,
            /*unconditional_features_file=*/nullptr);
   FUZZTEST_LOG(INFO) << "Number of input seeds available: "
                      << num_seeds_available
@@ -838,21 +844,39 @@ void Centipede::FuzzingLoop() {
     auto remaining_runs = env_.num_runs - new_runs;
     auto batch_size = std::min(env_.batch_size, remaining_runs);
     std::vector<MutationInputRef> mutation_inputs;
+    std::vector<size_t> mutate_batch_origins;
     mutation_inputs.reserve(env_.mutate_batch_size);
+    mutate_batch_origins.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 size_t origin = env_.use_corpus_weights
+                                ? corpus_.WeightedRandom(rng_)
+                                : corpus_.UniformRandom(rng_);
+      mutate_batch_origins.push_back(origin);
+      const auto& corpus_record = corpus_.Records()[origin];
       mutation_inputs.push_back(
           MutationInputRef{corpus_record.data, &corpus_record.metadata});
     }
 
-    const std::vector<ByteArray> mutants =
+    const std::vector<Mutant> mutants =
         user_callbacks_.Mutate(mutation_inputs, batch_size);
     if (ShouldStop()) break;
 
+    std::vector<ByteArray> next_batch;
+    next_batch.reserve(mutants.size());
+    std::vector<size_t> mutant_origins;
+    mutant_origins.reserve(mutants.size());
+    for (auto& mutant : mutants) {
+      next_batch.push_back(std::move(mutant.data));
+      if (mutant.origin == Mutant::kOriginNone) {
+        mutant_origins.push_back(Mutant::kOriginNone);
+      } else {
+        mutant_origins.push_back(mutate_batch_origins[mutant.origin]);
+      }
+    }
+
     bool gained_new_coverage =
-        RunBatch(mutants, corpus_file.get(), features_file.get(), nullptr);
+        RunBatch(next_batch, mutant_origins, corpus_file.get(),
+                 features_file.get(), nullptr);
     new_runs += mutants.size();
 
     if (gained_new_coverage) {

centipede/centipede.h

@@ -84,10 +84,11 @@ class Centipede {
   //   * its features are written to `features_file` (if that's non-null).
   // Returns true if new features were observed.
   // Post-condition: `batch_result.results.size()` == `input_vec.size()`.
-  bool RunBatch(const std::vector<ByteArray> &input_vec,
-                BlobFileWriter *absl_nullable corpus_file,
-                BlobFileWriter *absl_nullable features_file,
-                BlobFileWriter *absl_nullable unconditional_features_file);
+  bool RunBatch(const std::vector<ByteArray>& input_vec,
+                const std::vector<size_t>& mutant_origins,
+                BlobFileWriter* absl_nullable corpus_file,
+                BlobFileWriter* absl_nullable features_file,
+                BlobFileWriter* absl_nullable unconditional_features_file);
   // Loads seed inputs from the user callbacks, execute them, and store them
   // with the corresponding features into `corpus_file` and `features_file`.
   void LoadSeedInputs(BlobFileWriter *absl_nonnull corpus_file,

centipede/centipede_callbacks.h

@@ -72,8 +72,8 @@ class CentipedeCallbacks {
                        BatchResult &batch_result) = 0;
 
   // Takes non-empty `inputs` and returns at most `num_mutants` mutated inputs.
-  virtual std::vector<ByteArray> Mutate(
-      const std::vector<MutationInputRef> &inputs, size_t num_mutants) {
+  virtual std::vector<Mutant> Mutate(
+      const std::vector<MutationInputRef>& inputs, size_t num_mutants) {
     return env_.use_legacy_default_mutator
                ? byte_array_mutator_.MutateMany(inputs, num_mutants)
                : fuzztest_mutator_.MutateMany(inputs, num_mutants);

centipede/centipede_default_callbacks.cc

@@ -72,8 +72,8 @@ CentipedeDefaultCallbacks::GetSerializedTargetConfig() {
       "Failed to get serialized configuration from the target binary.");
 }
 
-std::vector<ByteArray> CentipedeDefaultCallbacks::Mutate(
-    const std::vector<MutationInputRef> &inputs, size_t num_mutants) {
+std::vector<Mutant> CentipedeDefaultCallbacks::Mutate(
+    const std::vector<MutationInputRef>& inputs, size_t num_mutants) {
   if (num_mutants == 0) return {};
   // Try to use the custom mutator if it hasn't been disabled.
   if (custom_mutator_is_usable_.value_or(true)) {

centipede/centipede_default_callbacks.h

@@ -42,8 +42,8 @@ class CentipedeDefaultCallbacks : public CentipedeCallbacks {
   absl::StatusOr<std::string> GetSerializedTargetConfig() override;
   bool Execute(std::string_view binary, const std::vector<ByteArray> &inputs,
                BatchResult &batch_result) override;
-  std::vector<ByteArray> Mutate(const std::vector<MutationInputRef> &inputs,
-                                size_t num_mutants) override;
+  std::vector<Mutant> Mutate(const std::vector<MutationInputRef>& inputs,
+                             size_t num_mutants) override;
 
  private:
   std::optional<bool> custom_mutator_is_usable_ = std::nullopt;

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(
+    bool, exec_time_weight_scaling, true,
+    "If true, scale the corpus weight by the execution time of each input.")
 CENTIPEDE_FLAG(
     bool, use_coverage_frontier, false,
     "If true, use coverage frontier when choosing the corpus element to "