From a617e16ff128c8aeac28212b5f06df212f6bbf3e Mon Sep 17 00:00:00 2001
From: Kostya Serebryany <kcc@google.com>
Date: Fri, 31 Mar 2017 04:17:45 +0000
Subject: [PATCH] [libFuzzer] simplify the code a bit

llvm-svn: 299180
---
 llvm/lib/Fuzzer/FuzzerCorpus.h  |  7 ++++---
 llvm/lib/Fuzzer/FuzzerLoop.cpp  | 12 ++++++------
 llvm/lib/Fuzzer/FuzzerTracePC.h | 12 ++++--------
 3 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/llvm/lib/Fuzzer/FuzzerCorpus.h b/llvm/lib/Fuzzer/FuzzerCorpus.h
index b6cbf36..0f05739 100644
--- a/llvm/lib/Fuzzer/FuzzerCorpus.h
+++ b/llvm/lib/Fuzzer/FuzzerCorpus.h
@@ -133,7 +133,7 @@ class InputCorpus {
       Printf("EVICTED %zd\n", Idx);
   }
 
-  bool AddFeature(size_t Idx, uint32_t NewSize, bool Shrink) {
+  void AddFeature(size_t Idx, uint32_t NewSize, bool Shrink) {
     assert(NewSize);
     Idx = Idx % kFeatureSetSize;
     uint32_t OldSize = GetFeature(Idx);
@@ -148,17 +148,17 @@ class InputCorpus {
       } else {
         NumAddedFeatures++;
       }
+      NumUpdatedFeatures++;
       if (FeatureDebug)
         Printf("ADD FEATURE %zd sz %d\n", Idx, NewSize);
       SmallestElementPerFeature[Idx] = Inputs.size();
       InputSizesPerFeature[Idx] = NewSize;
       CountingFeatures = true;
-      return true;
     }
-    return false;
   }
 
   size_t NumFeatures() const { return NumAddedFeatures; }
+  size_t NumFeatureUpdates() const { return NumUpdatedFeatures; }
 
   void ResetFeatureSet() {
     assert(Inputs.empty());
@@ -212,6 +212,7 @@ private:
 
   bool CountingFeatures = false;
   size_t NumAddedFeatures = 0;
+  size_t NumUpdatedFeatures = 0;
   uint32_t InputSizesPerFeature[kFeatureSetSize];
   uint32_t SmallestElementPerFeature[kFeatureSetSize];
 
diff --git a/llvm/lib/Fuzzer/FuzzerLoop.cpp b/llvm/lib/Fuzzer/FuzzerLoop.cpp
index 5c9dbc0..7040928 100644
--- a/llvm/lib/Fuzzer/FuzzerLoop.cpp
+++ b/llvm/lib/Fuzzer/FuzzerLoop.cpp
@@ -407,11 +407,11 @@ size_t Fuzzer::RunOne(const uint8_t *Data, size_t Size) {
 
   ExecuteCallback(Data, Size);
 
-  size_t Res = 0;
-  if (size_t NumFeatures = TPC.CollectFeatures([&](size_t Feature) -> bool {
-        return Corpus.AddFeature(Feature, Size, Options.Shrink);
-      }))
-    Res = NumFeatures;
+  size_t NumUpdatesBefore = Corpus.NumFeatureUpdates();
+  TPC.CollectFeatures([&](size_t Feature) {
+    Corpus.AddFeature(Feature, Size, Options.Shrink);
+  });
+  size_t NumUpdatesAfter = Corpus.NumFeatureUpdates();
 
   auto TimeOfUnit =
       duration_cast<seconds>(UnitStopTime - UnitStartTime).count();
@@ -424,7 +424,7 @@ size_t Fuzzer::RunOne(const uint8_t *Data, size_t Size) {
     Printf("Slowest unit: %zd s:\n", TimeOfLongestUnitInSeconds);
     WriteUnitToFileWithPrefix({Data, Data + Size}, "slow-unit-");
   }
-  return Res;
+  return NumUpdatesAfter - NumUpdatesBefore;
 }
 
 size_t Fuzzer::GetCurrentUnitInFuzzingThead(const uint8_t **Data) const {
diff --git a/llvm/lib/Fuzzer/FuzzerTracePC.h b/llvm/lib/Fuzzer/FuzzerTracePC.h
index 85c8e3c..6523fa0 100644
--- a/llvm/lib/Fuzzer/FuzzerTracePC.h
+++ b/llvm/lib/Fuzzer/FuzzerTracePC.h
@@ -58,7 +58,7 @@ class TracePC {
   void SetUseCounters(bool UC) { UseCounters = UC; }
   void SetUseValueProfile(bool VP) { UseValueProfile = VP; }
   void SetPrintNewPCs(bool P) { DoPrintNewPCs = P; }
-  template <class Callback> size_t CollectFeatures(Callback CB) const;
+  template <class Callback> void CollectFeatures(Callback CB) const;
 
   void ResetMaps() {
     ValueProfileMap.Reset();
@@ -129,8 +129,7 @@ void ForEachNonZeroByte(const uint8_t *Begin, const uint8_t *End,
 template <class Callback>  // bool Callback(size_t Feature)
 ATTRIBUTE_NO_SANITIZE_ALL
 __attribute__((noinline))
-size_t TracePC::CollectFeatures(Callback HandleFeature) const {
-  size_t Res = 0;
+void TracePC::CollectFeatures(Callback HandleFeature) const {
   uint8_t *Counters = this->Counters();
   size_t N = GetNumPCs();
   auto Handle8bitCounter = [&](size_t Idx, uint8_t Counter) {
@@ -143,8 +142,7 @@ size_t TracePC::CollectFeatures(Callback HandleFeature) const {
     else if (Counter >= 4) Bit = 3;
     else if (Counter >= 3) Bit = 2;
     else if (Counter >= 2) Bit = 1;
-    if (HandleFeature(Idx * 8 + Bit))
-      Res++;
+    HandleFeature(Idx * 8 + Bit);
   };
 
   ForEachNonZeroByte(Counters, Counters + N, 0, Handle8bitCounter);
@@ -153,10 +151,8 @@ size_t TracePC::CollectFeatures(Callback HandleFeature) const {
 
   if (UseValueProfile)
     ValueProfileMap.ForEach([&](size_t Idx) {
-      if (HandleFeature(N * 8 + Idx))
-        Res++;
+      HandleFeature(N * 8 + Idx);
     });
-  return Res;
 }
 
 extern TracePC TPC;
-- 
2.7.4