From a5b2e54fcb7d271ba2d92eda1d9c7cb9a4c4d6da Mon Sep 17 00:00:00 2001
From: Kostya Serebryany <kcc@google.com>
Date: Wed, 26 Oct 2016 00:20:51 +0000
Subject: [PATCH] [libFuzzer] simplify the code to print new PCs

llvm-svn: 285145
---
 llvm/lib/Fuzzer/FuzzerInternal.h  |  2 --
 llvm/lib/Fuzzer/FuzzerLoop.cpp    | 16 ++--------------
 llvm/lib/Fuzzer/FuzzerTracePC.cpp | 10 ++++++++++
 llvm/lib/Fuzzer/FuzzerTracePC.h   |  8 ++++++++
 4 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/Fuzzer/FuzzerInternal.h b/llvm/lib/Fuzzer/FuzzerInternal.h
index 15bf3e2..1b491ea 100644
--- a/llvm/lib/Fuzzer/FuzzerInternal.h
+++ b/llvm/lib/Fuzzer/FuzzerInternal.h
@@ -110,8 +110,6 @@ private:
   void InterruptCallback();
   void MutateAndTestOne();
   void ReportNewCoverage(InputInfo *II, const Unit &U);
-  void PrintNewPCs();
-  void PrintOneNewPC(uintptr_t PC);
   size_t RunOne(const Unit &U) { return RunOne(U.data(), U.size()); }
   void WriteToOutputCorpus(const Unit &U);
   void WriteUnitToFileWithPrefix(const Unit &U, const char *Prefix);
diff --git a/llvm/lib/Fuzzer/FuzzerLoop.cpp b/llvm/lib/Fuzzer/FuzzerLoop.cpp
index ad2e6d7..0c58728 100644
--- a/llvm/lib/Fuzzer/FuzzerLoop.cpp
+++ b/llvm/lib/Fuzzer/FuzzerLoop.cpp
@@ -168,6 +168,7 @@ Fuzzer::Fuzzer(UserCallback CB, InputCorpus &Corpus, MutationDispatcher &MD,
     EF->__sanitizer_install_malloc_and_free_hooks(MallocHook, FreeHook);
   TPC.SetUseCounters(Options.UseCounters);
   TPC.SetUseValueProfile(Options.UseValueProfile);
+  TPC.SetPrintNewPCs(Options.PrintNewCovPcs);
 
   if (Options.Verbosity)
     TPC.PrintModuleInfo();
@@ -556,26 +557,13 @@ void Fuzzer::PrintStatusForNewUnit(const Unit &U) {
   }
 }
 
-void Fuzzer::PrintOneNewPC(uintptr_t PC) {
-  PrintPC("\tNEW_PC: %p %F %L\n",
-          "\tNEW_PC: %p\n", PC);
-}
-
-void Fuzzer::PrintNewPCs() {
-  if (!Options.PrintNewCovPcs) return;
-  uintptr_t *PCIDs;
-  if (size_t NumNewPCIDs = TPC.GetNewPCIDs(&PCIDs))
-    for (size_t i = 0; i < NumNewPCIDs; i++)
-      PrintOneNewPC(TPC.GetPCbyPCID(PCIDs[i]));
-}
-
 void Fuzzer::ReportNewCoverage(InputInfo *II, const Unit &U) {
   II->NumSuccessfullMutations++;
   MD.RecordSuccessfulMutationSequence();
   PrintStatusForNewUnit(U);
   WriteToOutputCorpus(U);
   NumberOfNewUnitsAdded++;
-  PrintNewPCs();
+  TPC.PrintNewPCs();
 }
 
 // Finds minimal number of units in 'Extra' that add coverage to 'Initial'.
diff --git a/llvm/lib/Fuzzer/FuzzerTracePC.cpp b/llvm/lib/Fuzzer/FuzzerTracePC.cpp
index e17a10b..9e61247 100644
--- a/llvm/lib/Fuzzer/FuzzerTracePC.cpp
+++ b/llvm/lib/Fuzzer/FuzzerTracePC.cpp
@@ -136,6 +136,16 @@ static bool IsInterestingCoverageFile(std::string &File) {
   return true;
 }
 
+void TracePC::PrintNewPCs() {
+  if (DoPrintNewPCs) {
+    if (!PrintedPCs)
+      PrintedPCs = new std::set<uintptr_t>;
+    for (size_t i = 0; i < Min(NumGuards + 1, kNumPCs); i++)
+      if (PCs[i] && PrintedPCs->insert(PCs[i]).second)
+        PrintPC("\tNEW_PC: %p %F %L\n", "\tNEW_PC: %p\n", PCs[i]);
+  }
+}
+
 void TracePC::PrintCoverage() {
   if (!EF->__sanitizer_symbolize_pc) {
     Printf("INFO: __sanitizer_symbolize_pc is not available,"
diff --git a/llvm/lib/Fuzzer/FuzzerTracePC.h b/llvm/lib/Fuzzer/FuzzerTracePC.h
index 17c5575..06db3d0 100644
--- a/llvm/lib/Fuzzer/FuzzerTracePC.h
+++ b/llvm/lib/Fuzzer/FuzzerTracePC.h
@@ -12,6 +12,8 @@
 #ifndef LLVM_FUZZER_TRACE_PC
 #define LLVM_FUZZER_TRACE_PC
 
+#include <set>
+
 #include "FuzzerDefs.h"
 #include "FuzzerValueBitMap.h"
 
@@ -53,6 +55,7 @@ class TracePC {
   size_t GetTotalPCCoverage();
   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);
   bool UpdateValueProfileMap(ValueBitMap *MaxValueProfileMap) {
     return UseValueProfile && MaxValueProfileMap->MergeFrom(ValueProfileMap);
@@ -91,9 +94,12 @@ class TracePC {
   TableOfRecentCompares<uint32_t, kTORCSize> TORC4;
   TableOfRecentCompares<uint64_t, kTORCSize> TORC8;
 
+  void PrintNewPCs();
+
 private:
   bool UseCounters = false;
   bool UseValueProfile = false;
+  bool DoPrintNewPCs = false;
 
   static const size_t kMaxNewPCIDs = 1024;
   uintptr_t NewPCIDs[kMaxNewPCIDs];
@@ -129,6 +135,8 @@ private:
   static const size_t kNumPCs = 1 << 24;
   uintptr_t PCs[kNumPCs];
 
+  std::set<uintptr_t> *PrintedPCs;
+
   ValueBitMap ValueProfileMap;
 };
 
-- 
2.7.4