Create features file for distilled corpus (--distill flag) PiperOrigin-RevId: 561458581
diff --git a/centipede/BUILD b/centipede/BUILD index 3c5b9ca..52e507c 100644 --- a/centipede/BUILD +++ b/centipede/BUILD
@@ -1056,7 +1056,7 @@ ":distill", ":environment", ":feature", - ":logging", + ":shard_reader", ":test_util", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:reflection",
diff --git a/centipede/README.md b/centipede/README.md index 3ecda9d..82fadd0 100644 --- a/centipede/README.md +++ b/centipede/README.md
@@ -306,10 +306,10 @@ Then, run the same command line, but with `--distill --total_shards=N --num_threads=K`. This will read `N` corpus shards and produce `K` independent -distilled corpus files. Each of the distilled corpora should have the same -features as the `N` shards combined, but the inputs might be different between -the `K` distilled corpora. In most cases `K==1` is sufficient, i.e. you simply -omit `--num_threads=K`. +distilled corpus files and `K` corresponding feature files. Each of the +distilled corpora should have the same features as the `N` shards combined, but +the inputs might be different between the `K` distilled corpora. In most cases +`K==1` is sufficient, i.e. you simply omit `--num_threads=K`. The `--distill` flag requires that you pass the `--binary` or `--coverage_binary` so that it knows where to look for the `features` files, but @@ -345,6 +345,9 @@ ``` ... +├── distilled-features-byte_cmp_4.000000 +├── distilled-features-byte_cmp_4.000001 +├── distilled-features-byte_cmp_4.000002 ├── distilled-byte_cmp_4.000000 ├── distilled-byte_cmp_4.000001 ├── distilled-byte_cmp_4.000002
diff --git a/centipede/centipede.cc b/centipede/centipede.cc index d46fada..e21a673 100644 --- a/centipede/centipede.cc +++ b/centipede/centipede.cc
@@ -602,7 +602,7 @@ // Save the distilled corpus to a file in workdir and possibly to a hashed // file in the first corpus dir passed in `--corpus_dir`. - const auto distill_to_path = env_.MakeDistilledPath(); + const auto distill_to_path = env_.MakeDistilledCorpusPath(); LOG(INFO) << "Distilling: shard: " << env_.my_shard_index << " output: " << distill_to_path << " " << " distilled size: " << corpus_.NumActive();
diff --git a/centipede/distill.cc b/centipede/distill.cc index 85129de..95c35d1 100644 --- a/centipede/distill.cc +++ b/centipede/distill.cc
@@ -40,12 +40,17 @@ void DistillTask(const Environment &env, const std::vector<size_t> &shard_indices) { std::string log_line = absl::StrCat("DISTILL[S.", env.my_shard_index, "]: "); - const auto distill_to_path = env.MakeDistilledPath(); - LOG(INFO) << log_line << VV(env.total_shards) << VV(distill_to_path); + const auto corpus_path = env.MakeDistilledCorpusPath(); + const auto features_path = env.MakeDistilledFeaturesPath(); + LOG(INFO) << log_line << VV(env.total_shards) << VV(corpus_path) + << VV(features_path); - const auto appender = DefaultBlobFileWriterFactory(); - // NOTE: Overwrite distilled corpus files -- do not append. - CHECK_OK(appender->Open(distill_to_path, "w")); + const auto corpus_writer = DefaultBlobFileWriterFactory(); + const auto features_writer = DefaultBlobFileWriterFactory(); + // NOTE: Overwrite distilled corpus and features files -- do not append. + CHECK_OK(corpus_writer->Open(corpus_path, "w")); + CHECK_OK(features_writer->Open(features_path, "w")); + FeatureSet feature_set(/*frequency_threshold=*/1); for (size_t shard_idx : shard_indices) { LOG(INFO) << log_line << "reading shard " << shard_idx; @@ -66,7 +71,7 @@ // This is a simple linear greedy set cover algorithm. for (auto &&[input, features] : records) { VLOG(1) << log_line << VV(input.size()) << VV(features.size()); - if (!feature_set.CountUnseenAndPruneFrequentFeatures(features)) continue; + if (!feature_set.HasUnseenFeatures(features)) continue; feature_set.IncrementFrequencies(features); // Logging will log names of these variables. auto num_new_features = features.size(); @@ -75,8 +80,9 @@ auto ft = feature_set.size(); LOG(INFO) << log_line << "adding to distilled: " << VV(ft) << VV(cov) << VV(input.size()) << VV(num_new_features); - // Append to the distilled corpus. - CHECK_OK(appender->Write(input)); + // Append to the distilled corpus and features files. + CHECK_OK(corpus_writer->Write(input)); + CHECK_OK(features_writer->Write(PackFeaturesAndHash(input, features))); } } }
diff --git a/centipede/distill_test.cc b/centipede/distill_test.cc index b45c1fb..9d5c9cd 100644 --- a/centipede/distill_test.cc +++ b/centipede/distill_test.cc
@@ -25,6 +25,7 @@ #include "./centipede/defs.h" #include "./centipede/environment.h" #include "./centipede/feature.h" +#include "./centipede/shard_reader.h" #include "./centipede/test_util.h" ABSL_DECLARE_FLAG(std::string, binary_hash); @@ -39,6 +40,19 @@ FeatureVec feature_vec; }; +// Custom matcher for TestCorpusRecord. Compares `expected_input` with +// actual TestCorpusRecord::input and compares `expected_features` with +// actual TestCorpusRecord::feature_vec. +MATCHER_P2(EqualsTestCorpusRecord, expected_input, expected_features, "") { + return testing::ExplainMatchResult( + testing::Field(&TestCorpusRecord::input, expected_input), arg, + result_listener) && + testing::ExplainMatchResult( + testing::Field(&TestCorpusRecord::feature_vec, + testing::ElementsAreArray(expected_features)), + arg, result_listener); +} + using Shard = std::vector<TestCorpusRecord>; using ShardVec = std::vector<Shard>; using InputVec = std::vector<ByteArray>; @@ -57,24 +71,27 @@ PackFeaturesAndHash(record.input, record.feature_vec))); } -// Reads and returns the distilled corpus from `env.MakeDistilledPath()`. -std::vector<ByteArray> ReadFromDistilled(const Environment &env) { - auto distilled_path = env.MakeDistilledPath(); - auto reader = DefaultBlobFileReaderFactory(); - CHECK_OK(reader->Open(distilled_path)); - absl::Span<uint8_t> blob; - std::vector<ByteArray> result; - while (reader->Read(blob).ok()) { - result.emplace_back(blob.begin(), blob.end()); - } +// Reads and returns the distilled corpus record from +// `env.MakeDistilledCorpusPath()` and `env.MakeDistilledFeaturesPath()`. +std::vector<TestCorpusRecord> ReadFromDistilled(const Environment &env) { + auto distilled_corpus_path = env.MakeDistilledCorpusPath(); + auto distilled_features_path = env.MakeDistilledFeaturesPath(); + + std::vector<TestCorpusRecord> result; + auto shard_reader_callback = [&result](const ByteArray &input, + FeatureVec &features) { + result.push_back({input, features}); + }; + ReadShard(env.MakeDistilledCorpusPath(), env.MakeDistilledFeaturesPath(), + shard_reader_callback); return result; } // Distills `shards` in the order specified by `shard_indices`, // returns the distilled corpus as a vector of inputs. -InputVec TestDistill(const ShardVec &shards, - const std::vector<size_t> &shard_indices, - std::string_view test_name) { +std::vector<TestCorpusRecord> TestDistill( + const ShardVec &shards, const std::vector<size_t> &shard_indices, + std::string_view test_name) { // Set up the environment. // We need to set at least --binary_hash before `env` is constructed, // so we do this by overriding the flags. @@ -117,11 +134,22 @@ }; // Distill these 3 shards in different orders, observe different results. EXPECT_THAT(TestDistill(shards, {0, 1, 2}, test_info_->name()), - testing::ElementsAreArray({in0, in1, in2})); + testing::ElementsAreArray({ + EqualsTestCorpusRecord(in0, FeatureVec{10, 20}), + EqualsTestCorpusRecord(in1, FeatureVec{20, 30}), + EqualsTestCorpusRecord(in2, FeatureVec{30, 40}), + })); EXPECT_THAT(TestDistill(shards, {2, 0, 1}, test_info_->name()), - testing::ElementsAreArray({in2, in0})); + testing::ElementsAreArray({ + EqualsTestCorpusRecord(in2, FeatureVec{30, 40}), + EqualsTestCorpusRecord(in0, FeatureVec{10, 20}), + })); EXPECT_THAT(TestDistill(shards, {1, 0, 2}, test_info_->name()), - testing::ElementsAreArray({in1, in0, in2})); + testing::ElementsAreArray({ + EqualsTestCorpusRecord(in1, FeatureVec{20, 30}), + EqualsTestCorpusRecord(in0, FeatureVec{10, 20}), + EqualsTestCorpusRecord(in2, FeatureVec{30, 40}), + })); } // TODO(kcc): add more tests once we settle on the testing code above.
diff --git a/centipede/environment.cc b/centipede/environment.cc index 753c621..c13c1f3 100644 --- a/centipede/environment.cc +++ b/centipede/environment.cc
@@ -552,11 +552,17 @@ absl::StrFormat("features.%0*d", kDigitsInShardIndex, shard_index)); } -std::string Environment::MakeDistilledPath() const { +std::string Environment::MakeDistilledCorpusPath() const { return std::filesystem::path(workdir).append(absl::StrFormat( "distilled-%s.%0*d", binary_name, kDigitsInShardIndex, my_shard_index)); } +std::string Environment::MakeDistilledFeaturesPath() const { + return std::filesystem::path(workdir).append( + absl::StrFormat("distilled-features-%s.%0*d", binary_name, + kDigitsInShardIndex, my_shard_index)); +} + std::string Environment::MakeCoverageReportPath( std::string_view annotation) const { return std::filesystem::path(workdir).append(absl::StrFormat(
diff --git a/centipede/environment.h b/centipede/environment.h index 9d36950..db92133 100644 --- a/centipede/environment.h +++ b/centipede/environment.h
@@ -154,7 +154,9 @@ // Returns the path to the indexed code coverage file. std::string MakeSourceBasedCoverageIndexedProfilePath() const; // Returns the path for the distilled corpus file for my_shard_index. - std::string MakeDistilledPath() const; + std::string MakeDistilledCorpusPath() const; + // Returns the path for the distilled features file for my_shard_index. + std::string MakeDistilledFeaturesPath() const; // Returns the path for the coverage report file for my_shard_index. // Non-default `annotation` becomes a part of the returned filename. // `annotation` must not start with a '.'.
diff --git a/centipede/environment_test.cc b/centipede/environment_test.cc index d035890..d6f384f 100644 --- a/centipede/environment_test.cc +++ b/centipede/environment_test.cc
@@ -49,6 +49,13 @@ Experiment(11, true, 30, "E12", "use_cmp_features=true:path_level=30:"); } +TEST(Environment, MakeDistilledCorpusAndFeaturesPaths) { + Environment env; + env.my_shard_index = 3; + EXPECT_EQ(env.MakeDistilledCorpusPath(), "distilled-.000003"); + EXPECT_EQ(env.MakeDistilledFeaturesPath(), "distilled-features-.000003"); +} + TEST(Environment, MakeCoverageReportPath) { // TODO(ussuri): Environment is not test-friendly (initialized through // flags, which are hidden in the .cc). Fix.
diff --git a/centipede/feature_set.cc b/centipede/feature_set.cc index aa93a35..d04dac7 100644 --- a/centipede/feature_set.cc +++ b/centipede/feature_set.cc
@@ -40,13 +40,19 @@ return features_per_domain_[domain.domain_id()]; } +bool FeatureSet::HasUnseenFeatures(const FeatureVec &features) const { + for (auto feature : features) { + if (frequencies_[feature] == 0) return true; + } + return false; +} + __attribute__((noinline)) // to see it in profile. size_t FeatureSet::CountUnseenAndPruneFrequentFeatures(FeatureVec &features) const { size_t number_of_unseen_features = 0; size_t num_kept = 0; - for (size_t i = 0, n = features.size(); i < n; i++) { - auto feature = features[i]; + for (auto feature : features) { auto freq = frequencies_[feature]; if (freq == 0) { ++number_of_unseen_features;
diff --git a/centipede/feature_set.h b/centipede/feature_set.h index 8032fa9..4daddf9 100644 --- a/centipede/feature_set.h +++ b/centipede/feature_set.h
@@ -37,6 +37,9 @@ explicit FeatureSet(uint8_t frequency_threshold) : frequency_threshold_(frequency_threshold) {} + // Returns true if there are features in `features` not present in `this`. + bool HasUnseenFeatures(const FeatureVec &features) const; + // Returns the number of features in `features` not present in `this`. // Removes all features from `features` that are too frequent. size_t CountUnseenAndPruneFrequentFeatures(FeatureVec &features) const;
diff --git a/centipede/feature_set_test.cc b/centipede/feature_set_test.cc index 60d9144..9b89ca1 100644 --- a/centipede/feature_set_test.cc +++ b/centipede/feature_set_test.cc
@@ -68,6 +68,28 @@ EXPECT_GT(weight({f2}), weight({f3})); } +TEST(FeatureSet, HasUnseenFeatures_IncrementFrequencies) { + size_t frequency_threshold = 2; + FeatureSet feature_set(frequency_threshold); + FeatureVec features = {10}; + EXPECT_TRUE(feature_set.HasUnseenFeatures(features)); + + feature_set.IncrementFrequencies(features); + EXPECT_FALSE(feature_set.HasUnseenFeatures(features)); + + features = {10, 20}; + EXPECT_TRUE(feature_set.HasUnseenFeatures(features)); + feature_set.IncrementFrequencies(features); + EXPECT_FALSE(feature_set.HasUnseenFeatures(features)); + + features = {50}; + EXPECT_TRUE(feature_set.HasUnseenFeatures(features)); + feature_set.IncrementFrequencies(features); + + features = {10, 20}; + EXPECT_FALSE(feature_set.HasUnseenFeatures(features)); +} + TEST(FeatureSet, CountUnseenAndPruneFrequentFeatures_IncrementFrequencies) { size_t frequency_threshold = 3; FeatureSet feature_set(frequency_threshold);