2 files changed, 34 insertions, 8 deletions

diff --git a/compiler-rt/lib/fuzzer/FuzzerFork.cpp b/compiler-rt/lib/fuzzer/FuzzerFork.cpp
index 20c9950b349..9f315e8c125 100644
--- a/compiler-rt/lib/fuzzer/FuzzerFork.cpp
+++ b/compiler-rt/lib/fuzzer/FuzzerFork.cpp
@@ -11,6 +11,7 @@
 #include "FuzzerCommand.h"
 #include "FuzzerFork.h"
 #include "FuzzerIO.h"
+#include "FuzzerInternal.h"
 #include "FuzzerMerge.h"
 #include "FuzzerSHA1.h"
 #include "FuzzerTracePC.h"
@@ -63,6 +64,7 @@ struct FuzzJob {
   // Inputs.
   Command Cmd;
   std::string CorpusDir;
+  std::string FeaturesDir;
   std::string LogPath;
   std::string CFPath;
 
@@ -73,6 +75,7 @@ struct FuzzJob {
     RemoveFile(CFPath);
     RemoveFile(LogPath);
     RmDirRecursive(CorpusDir);
+    RmDirRecursive(FeaturesDir);
   }
 };
 
@@ -122,12 +125,17 @@ struct GlobalEnv {
       Cmd.addFlag("seed_inputs", Seeds);
     Job->LogPath = DirPlusFile(TempDir, std::to_string(JobId) + ".log");
     Job->CorpusDir = DirPlusFile(TempDir, "C" + std::to_string(JobId));
+    Job->FeaturesDir = DirPlusFile(TempDir, "F" + std::to_string(JobId));
     Job->CFPath = DirPlusFile(TempDir, std::to_string(JobId) + ".merge");
 
 
     Cmd.addArgument(Job->CorpusDir);
-    RmDirRecursive(Job->CorpusDir);
-    MkDir(Job->CorpusDir);
+    Cmd.addFlag("features_dir", Job->FeaturesDir);
+
+    for (auto &D : {Job->CorpusDir, Job->FeaturesDir}) {
+      RmDirRecursive(D);
+      MkDir(D);
+    }
 
     Cmd.setOutputFile(Job->LogPath);
     Cmd.combineOutAndErr();
@@ -142,12 +150,30 @@ struct GlobalEnv {
   }
 
   void RunOneMergeJob(FuzzJob *Job) {
-    Vector<SizedFile> TempFiles;
+    Vector<SizedFile> TempFiles, MergeCandidates;
+    // Read all newly created inputs and their feature sets.
+    // Choose only those inputs that have new features.
     GetSizedFilesFromDir(Job->CorpusDir, &TempFiles);
+    std::sort(TempFiles.begin(), TempFiles.end());
+    for (auto &F : TempFiles) {
+      auto FeatureFile = F.File;
+      FeatureFile.replace(0, Job->CorpusDir.size(), Job->FeaturesDir);
+      auto FeatureBytes = FileToVector(FeatureFile, 0, false);
+      assert((FeatureBytes.size() % sizeof(uint32_t)) == 0);
+      Vector<uint32_t> NewFeatures(FeatureBytes.size() / sizeof(uint32_t));
+      memcpy(NewFeatures.data(), FeatureBytes.data(), FeatureBytes.size());
+      for (auto Ft : NewFeatures) {
+        if (!Features.count(Ft)) {
+          MergeCandidates.push_back(F);
+          break;
+        }
+      }
+    }
+    if (MergeCandidates.empty()) return;
 
     Vector<std::string> FilesToAdd;
     Set<uint32_t> NewFeatures, NewCov;
-    CrashResistantMerge(Args, {}, TempFiles, &FilesToAdd, Features,
+    CrashResistantMerge(Args, {}, MergeCandidates, &FilesToAdd, Features,
                         &NewFeatures, Cov, &NewCov, Job->CFPath, false);
     for (auto &Path : FilesToAdd) {
       auto U = FileToVector(Path);
@@ -265,6 +291,7 @@ void FuzzWithFork(Random &Rand, const FuzzingOptions &Options,
       Stop = true;
       break;
     }
+    Fuzzer::MaybeExitGracefully();
 
     Env.RunOneMergeJob(Job.get());
 
diff --git a/compiler-rt/lib/fuzzer/FuzzerLoop.cpp b/compiler-rt/lib/fuzzer/FuzzerLoop.cpp
index b68d759cf83..a323a7a465b 100644
--- a/compiler-rt/lib/fuzzer/FuzzerLoop.cpp
+++ b/compiler-rt/lib/fuzzer/FuzzerLoop.cpp
@@ -446,13 +446,12 @@ void Fuzzer::PrintPulseAndReportSlowInput(const uint8_t *Data, size_t Size) {
 }
 
 static void WriteFeatureSetToFile(const std::string &FeaturesDir,
-                                  const uint8_t Sha1[],
+                                  const std::string &FileName,
                                   const Vector<uint32_t> &FeatureSet) {
   if (FeaturesDir.empty() || FeatureSet.empty()) return;
   WriteToFile(reinterpret_cast<const uint8_t *>(FeatureSet.data()),
               FeatureSet.size() * sizeof(FeatureSet[0]),
-              DirPlusFile(FeaturesDir, Sha1ToString(Sha1)));
-  Printf("Features: %s\n", Sha1ToString(Sha1).c_str());
+              DirPlusFile(FeaturesDir, FileName));
 }
 
 static void RenameFeatureSetFile(const std::string &FeaturesDir,
@@ -490,7 +489,7 @@ bool Fuzzer::RunOne(const uint8_t *Data, size_t Size, bool MayDeleteFile,
     auto NewII = Corpus.AddToCorpus({Data, Data + Size}, NumNewFeatures,
                                     MayDeleteFile, TPC.ObservedFocusFunction(),
                                     UniqFeatureSetTmp, DFT, II);
-    WriteFeatureSetToFile(Options.FeaturesDir, NewII->Sha1,
+    WriteFeatureSetToFile(Options.FeaturesDir, Sha1ToString(NewII->Sha1),
                           NewII->UniqFeatureSet);
     return true;
   }