From fe1094b8115ea8ded42d8d2c3dbcfa9ef32f0a76 Mon Sep 17 00:00:00 2001
From: Kostya Serebryany <kcc@google.com>
Date: Mon, 5 Dec 2016 23:35:22 +0000
Subject: [PATCH] [libFuzzer] refactor the code to allow collecting features in
 different ways. Also initialize a couple of Fuzzer:: members that might have
 been used uninitialized :(

llvm-svn: 288731
---
 llvm/lib/Fuzzer/FuzzerLoop.cpp    |  6 +++++-
 llvm/lib/Fuzzer/FuzzerTracePC.cpp | 35 -----------------------------------
 llvm/lib/Fuzzer/FuzzerTracePC.h   | 38 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 42 insertions(+), 37 deletions(-)

diff --git a/llvm/lib/Fuzzer/FuzzerLoop.cpp b/llvm/lib/Fuzzer/FuzzerLoop.cpp
index e5f3af9..3d90401 100644
--- a/llvm/lib/Fuzzer/FuzzerLoop.cpp
+++ b/llvm/lib/Fuzzer/FuzzerLoop.cpp
@@ -193,6 +193,8 @@ Fuzzer::Fuzzer(UserCallback CB, InputCorpus &Corpus, MutationDispatcher &MD,
     EpochOfLastReadOfOutputCorpus = GetEpoch(Options.OutputCorpus);
   MaxInputLen = MaxMutationLen = Options.MaxLen;
   AllocateCurrentUnitData();
+  CurrentUnitSize = 0;
+  memset(BaseSha1, 0, sizeof(BaseSha1));
 }
 
 Fuzzer::~Fuzzer() { }
@@ -486,7 +488,9 @@ size_t Fuzzer::RunOne(const uint8_t *Data, size_t Size) {
   ExecuteCallback(Data, Size);
 
   size_t Res = 0;
-  if (size_t NumFeatures = TPC.FinalizeTrace(&Corpus, Size, Options.Shrink))
+  if (size_t NumFeatures = TPC.CollectFeatures([&](size_t Feature) -> bool {
+        return Corpus.AddFeature(Feature, Size, Options.Shrink);
+      }))
     Res = NumFeatures;
 
   if (!TPC.UsingTracePcGuard()) {
diff --git a/llvm/lib/Fuzzer/FuzzerTracePC.cpp b/llvm/lib/Fuzzer/FuzzerTracePC.cpp
index 8d58a6d..d8036ed 100644
--- a/llvm/lib/Fuzzer/FuzzerTracePC.cpp
+++ b/llvm/lib/Fuzzer/FuzzerTracePC.cpp
@@ -59,41 +59,6 @@ void TracePC::PrintModuleInfo() {
   Printf("\n");
 }
 
-size_t TracePC::FinalizeTrace(InputCorpus *C, size_t InputSize, bool Shrink) {
-  if (!UsingTracePcGuard()) return 0;
-  size_t Res = 0;
-  const size_t Step = 8;
-  assert(reinterpret_cast<uintptr_t>(Counters) % Step == 0);
-  size_t N = Min(kNumCounters, NumGuards + 1);
-  N = (N + Step - 1) & ~(Step - 1);  // Round up.
-  for (size_t Idx = 0; Idx < N; Idx += Step) {
-    uint64_t Bundle = *reinterpret_cast<uint64_t*>(&Counters[Idx]);
-    if (!Bundle) continue;
-    for (size_t i = Idx; i < Idx + Step; i++) {
-      uint8_t Counter = (Bundle >> (i * 8)) & 0xff;
-      if (!Counter) continue;
-      Counters[i] = 0;
-      unsigned Bit = 0;
-      /**/ if (Counter >= 128) Bit = 7;
-      else if (Counter >= 32) Bit = 6;
-      else if (Counter >= 16) Bit = 5;
-      else if (Counter >= 8) Bit = 4;
-      else if (Counter >= 4) Bit = 3;
-      else if (Counter >= 3) Bit = 2;
-      else if (Counter >= 2) Bit = 1;
-      size_t Feature = (i * 8 + Bit);
-      if (C->AddFeature(Feature, InputSize, Shrink))
-        Res++;
-    }
-  }
-  if (UseValueProfile)
-    ValueProfileMap.ForEach([&](size_t Idx) {
-      if (C->AddFeature(NumGuards + Idx, InputSize, Shrink))
-        Res++;
-    });
-  return Res;
-}
-
 void TracePC::HandleCallerCallee(uintptr_t Caller, uintptr_t Callee) {
   const uintptr_t kBits = 12;
   const uintptr_t kMask = (1 << kBits) - 1;
diff --git a/llvm/lib/Fuzzer/FuzzerTracePC.h b/llvm/lib/Fuzzer/FuzzerTracePC.h
index 9c7f563..acff27f 100644
--- a/llvm/lib/Fuzzer/FuzzerTracePC.h
+++ b/llvm/lib/Fuzzer/FuzzerTracePC.h
@@ -56,7 +56,7 @@ class TracePC {
   void SetUseCounters(bool UC) { UseCounters = UC; }
   void SetUseValueProfile(bool VP) { UseValueProfile = VP; }
   void SetPrintNewPCs(bool P) { DoPrintNewPCs = P; }
-  size_t FinalizeTrace(InputCorpus *C, size_t InputSize, bool Shrink);
+  template <class Callback> size_t CollectFeatures(Callback CB);
   bool UpdateValueProfileMap(ValueBitMap *MaxValueProfileMap) {
     return UseValueProfile && MaxValueProfileMap->MergeFrom(ValueProfileMap);
   }
@@ -115,6 +115,42 @@ private:
   ValueBitMap ValueProfileMap;
 };
 
+template <class Callback>
+size_t TracePC::CollectFeatures(Callback CB) {
+  if (!UsingTracePcGuard()) return 0;
+  size_t Res = 0;
+  const size_t Step = 8;
+  assert(reinterpret_cast<uintptr_t>(Counters) % Step == 0);
+  size_t N = Min(kNumCounters, NumGuards + 1);
+  N = (N + Step - 1) & ~(Step - 1);  // Round up.
+  for (size_t Idx = 0; Idx < N; Idx += Step) {
+    uint64_t Bundle = *reinterpret_cast<uint64_t*>(&Counters[Idx]);
+    if (!Bundle) continue;
+    for (size_t i = Idx; i < Idx + Step; i++) {
+      uint8_t Counter = (Bundle >> (i * 8)) & 0xff;
+      if (!Counter) continue;
+      Counters[i] = 0;
+      unsigned Bit = 0;
+      /**/ if (Counter >= 128) Bit = 7;
+      else if (Counter >= 32) Bit = 6;
+      else if (Counter >= 16) Bit = 5;
+      else if (Counter >= 8) Bit = 4;
+      else if (Counter >= 4) Bit = 3;
+      else if (Counter >= 3) Bit = 2;
+      else if (Counter >= 2) Bit = 1;
+      size_t Feature = (i * 8 + Bit);
+      if (CB(Feature))
+        Res++;
+    }
+  }
+  if (UseValueProfile)
+    ValueProfileMap.ForEach([&](size_t Idx) {
+      if (CB(NumGuards + Idx))
+        Res++;
+    });
+  return Res;
+}
+
 extern TracePC TPC;
 
 }  // namespace fuzzer
-- 
2.7.4