From 1e438a1dbae0c5c37aabe99becd279a0ccfb7566 Mon Sep 17 00:00:00 2001
From: Kostya Serebryany <kcc@google.com>
Date: Sat, 17 Dec 2016 08:20:24 +0000
Subject: [PATCH] [libFuzzer] use less memory for merge

llvm-svn: 290039
---
 llvm/lib/Fuzzer/FuzzerMerge.cpp         | 13 +++++++------
 llvm/lib/Fuzzer/FuzzerMerge.h           |  2 +-
 llvm/lib/Fuzzer/test/FuzzerUnittest.cpp |  7 +++++--
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Fuzzer/FuzzerMerge.cpp b/llvm/lib/Fuzzer/FuzzerMerge.cpp
index 3847995..84660e0 100644
--- a/llvm/lib/Fuzzer/FuzzerMerge.cpp
+++ b/llvm/lib/Fuzzer/FuzzerMerge.cpp
@@ -94,10 +94,11 @@ bool Merger::Parse(std::istream &IS, bool ParseCoverage) {
         return false;
       LastSeenStartMarker = kInvalidStartMarker;
       if (ParseCoverage) {
-        while (!ISS1.rdstate()) {
-          ISS1 >> std::hex >> N;
-          Files[CurrentFileIdx].Features.insert(N);
-        }
+        auto &V = Files[CurrentFileIdx].Features;
+        V.clear();
+        while (ISS1 >> std::hex >> N)
+          V.push_back(N);
+        std::sort(V.begin(), V.end());
       }
     } else {
       return false;
@@ -115,7 +116,7 @@ bool Merger::Parse(std::istream &IS, bool ParseCoverage) {
 size_t Merger::Merge(std::vector<std::string> *NewFiles) {
   NewFiles->clear();
   assert(NumFilesInFirstCorpus <= Files.size());
-  std::set<size_t> AllFeatures;
+  std::set<uint32_t> AllFeatures;
 
   // What features are in the initial corpus?
   for (size_t i = 0; i < NumFilesInFirstCorpus; i++) {
@@ -127,7 +128,7 @@ size_t Merger::Merge(std::vector<std::string> *NewFiles) {
   // Remove all features that we already know from all other inputs.
   for (size_t i = NumFilesInFirstCorpus; i < Files.size(); i++) {
     auto &Cur = Files[i].Features;
-    std::set<size_t> Tmp;
+    std::vector<uint32_t> Tmp;
     std::set_difference(Cur.begin(), Cur.end(), AllFeatures.begin(),
                         AllFeatures.end(), std::inserter(Tmp, Tmp.begin()));
     Cur.swap(Tmp);
diff --git a/llvm/lib/Fuzzer/FuzzerMerge.h b/llvm/lib/Fuzzer/FuzzerMerge.h
index d00349a..8a2fe5d 100644
--- a/llvm/lib/Fuzzer/FuzzerMerge.h
+++ b/llvm/lib/Fuzzer/FuzzerMerge.h
@@ -50,7 +50,7 @@ namespace fuzzer {
 struct MergeFileInfo {
   std::string Name;
   size_t Size = 0;
-  std::set<size_t> Features;
+  std::vector<uint32_t> Features;
 };
 
 struct Merger {
diff --git a/llvm/lib/Fuzzer/test/FuzzerUnittest.cpp b/llvm/lib/Fuzzer/test/FuzzerUnittest.cpp
index aab32622..4992ef5 100644
--- a/llvm/lib/Fuzzer/test/FuzzerUnittest.cpp
+++ b/llvm/lib/Fuzzer/test/FuzzerUnittest.cpp
@@ -620,7 +620,7 @@ TEST(Merge, Bad) {
   }
 }
 
-void EQ(const std::set<size_t> &A, const std::set<size_t> &B) {
+void EQ(const std::vector<uint32_t> &A, const std::vector<uint32_t> &B) {
   EXPECT_EQ(A, B);
 }
 
@@ -699,8 +699,11 @@ TEST(Merge, Good) {
   EXPECT_TRUE(M.Parse("3\n1\nA\nB\nC\n"
                         "STARTED 0 1000\nDONE 0 1 2 3\n"
                         "STARTED 1 1001\nDONE 1 4 5 6 \n"
-                        "STARTED 2 1002\nDONE 2 6 1 3 \n"
+                        "STARTED 2 1002\nDONE 2 6 1 3\n"
                         "", true));
+  EQ(M.Files[0].Features, {1, 2, 3});
+  EQ(M.Files[1].Features, {4, 5, 6});
+  EQ(M.Files[2].Features, {1, 3, 6});
   EXPECT_EQ(3U, M.Merge(&NewFiles));
   EQ(NewFiles, {"B"});
 }
-- 
2.7.4