7 files changed, 43 insertions, 1 deletions
diff --git a/lib/Fuzzer/FuzzerDriver.cpp b/lib/Fuzzer/FuzzerDriver.cpp
index 1746afd..9ccd744 100644
--- a/lib/Fuzzer/FuzzerDriver.cpp
+++ b/lib/Fuzzer/FuzzerDriver.cpp
@@ -158,6 +158,7 @@ int FuzzerDriver(int argc, char **argv, UserCallback Callback) {
   Options.DoCrossOver = Flags.cross_over;
   Options.MutateDepth = Flags.mutate_depth;
   Options.ExitOnFirst = Flags.exit_on_first;
+  Options.UseCounters = Flags.use_counters;
   Options.UseFullCoverageSet = Flags.use_full_coverage_set;
   Options.UseCoveragePairs = Flags.use_coverage_pairs;
   Options.PreferSmallDuringInitialShuffle =
diff --git a/lib/Fuzzer/FuzzerFlags.def b/lib/Fuzzer/FuzzerFlags.def
index 068f245..08176af 100644
--- a/lib/Fuzzer/FuzzerFlags.def
+++ b/lib/Fuzzer/FuzzerFlags.def
@@ -32,6 +32,7 @@ FUZZER_FLAG(int, help, 0, "Print help.")
 FUZZER_FLAG(
     int, save_minimized_corpus, 0,
     "If 1, the minimized corpus is saved into the first input directory")
+FUZZER_FLAG(int, use_counters, 0, "Use coverage counters")
 FUZZER_FLAG(int, use_full_coverage_set, 0,
             "Experimental: Maximize the number of different full"
             " coverage sets as opposed to maximizing the total coverage."
diff --git a/lib/Fuzzer/FuzzerInternal.h b/lib/Fuzzer/FuzzerInternal.h
index 980b00e..e4e5eb7 100644
--- a/lib/Fuzzer/FuzzerInternal.h
+++ b/lib/Fuzzer/FuzzerInternal.h
@@ -48,6 +48,7 @@ class Fuzzer {
     bool DoCrossOver = true;
     int  MutateDepth = 5;
     bool ExitOnFirst = false;
+    bool UseCounters = false;
     bool UseFullCoverageSet  = false;
     bool UseCoveragePairs = false;
     int PreferSmallDuringInitialShuffle = -1;
@@ -95,6 +96,15 @@ class Fuzzer {
   std::vector<Unit> Corpus;
   std::unordered_set<uintptr_t> FullCoverageSets;
   std::unordered_set<uint64_t>  CoveragePairs;
+
+  // For UseCounters
+  std::vector<uint8_t> CounterBitmap;
+  size_t TotalBits() {  // Slow. Call it only for printing stats.
+    size_t Res = 0;
+    for (auto x : CounterBitmap) Res += __builtin_popcount(x);
+    return Res;
+  }
+
   UserCallback Callback;
   FuzzingOptions Options;
   system_clock::time_point ProcessStartTime = system_clock::now();
diff --git a/lib/Fuzzer/FuzzerLoop.cpp b/lib/Fuzzer/FuzzerLoop.cpp
index 70b63eb..563fbf4 100644
--- a/lib/Fuzzer/FuzzerLoop.cpp
+++ b/lib/Fuzzer/FuzzerLoop.cpp
@@ -138,17 +138,28 @@ size_t Fuzzer::RunOneMaximizeFullCoverageSet(const Unit &U) {
 }
 
 size_t Fuzzer::RunOneMaximizeTotalCoverage(const Unit &U) {
+  size_t NumCounters = __sanitizer_get_number_of_counters();
+  if (Options.UseCounters) {
+    CounterBitmap.resize(NumCounters);
+    __sanitizer_update_counter_bitset_and_clear_counters(0);
+  }
   size_t OldCoverage = __sanitizer_get_total_unique_coverage();
   Callback(U.data(), U.size());
   size_t NewCoverage = __sanitizer_get_total_unique_coverage();
+  size_t NumNewBits = 0;
+  if (Options.UseCounters)
+    NumNewBits = __sanitizer_update_counter_bitset_and_clear_counters(
+        CounterBitmap.data());
+
   if (!(TotalNumberOfRuns & (TotalNumberOfRuns - 1)) && Options.Verbosity) {
     size_t Seconds = secondsSinceProcessStartUp();
     std::cerr
         << "#" << TotalNumberOfRuns
         << "\tcov: " << NewCoverage
+        << "\tbits: " << TotalBits()
         << "\texec/s: " << (Seconds ? TotalNumberOfRuns / Seconds : 0) << "\n";
   }
-  if (NewCoverage > OldCoverage)
+  if (NewCoverage > OldCoverage || NumNewBits)
     return NewCoverage;
   return 0;
 }
@@ -189,6 +200,7 @@ size_t Fuzzer::MutateAndTestOne(Unit *U) {
       if (Options.Verbosity) {
         std::cerr << "#" << TotalNumberOfRuns
                   << "\tNEW: " << NewCoverage
+                  << " B: " << TotalBits()
                   << " L: " << U->size()
                   << " S: " << Corpus.size()
                   << " I: " << i
diff --git a/lib/Fuzzer/test/CMakeLists.txt b/lib/Fuzzer/test/CMakeLists.txt
index bed9cd8..08130c6 100644
--- a/lib/Fuzzer/test/CMakeLists.txt
+++ b/lib/Fuzzer/test/CMakeLists.txt
@@ -5,6 +5,7 @@
 set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -O0 -fsanitize-coverage=4")
 
 set(Tests
+  CounterTest
   FourIndependentBranchesTest
   FullCoverageSetTest
   InfiniteTest
diff --git a/lib/Fuzzer/test/CounterTest.cpp b/lib/Fuzzer/test/CounterTest.cpp
new file mode 100644
index 0000000..332ccfe
--- /dev/null
+++ b/lib/Fuzzer/test/CounterTest.cpp
@@ -0,0 +1,14 @@
+// Test for a fuzzer: must find the case where a particular basic block is
+// executed many times.
+#include <iostream>
+
+extern "C" void TestOneInput(const uint8_t *Data, size_t Size) {
+  int Num = 0;
+  for (size_t i = 0; i < Size; i++)
+    if (Data[i] == 'A' + i)
+      Num++;
+  if (Num >= 4) {
+    std::cerr <<  "BINGO!\n";
+    exit(1);
+  }
+}
diff --git a/lib/Fuzzer/test/fuzzer.test b/lib/Fuzzer/test/fuzzer.test
index 1e42e72..45691f5 100644
--- a/lib/Fuzzer/test/fuzzer.test
+++ b/lib/Fuzzer/test/fuzzer.test
@@ -17,3 +17,6 @@ FullCoverageSetTest: BINGO
 
 RUN: not ./LLVMFuzzer-FourIndependentBranchesTest -timeout=15 -seed=1 -use_coverage_pairs=1 2>&1 | FileCheck %s --check-prefix=FourIndependentBranchesTest
 FourIndependentBranchesTest: BINGO
+
+RUN: not ./LLVMFuzzer-CounterTest -use_counters=1 -max_len=6 -seed=1 -timeout=15 2>&1 | FileCheck %s --check-prefix=CounterTest
+CounterTest: BINGO