2 files changed, 73 insertions, 0 deletions
diff --git a/lib/sanitizer_common/sanitizer_coverage_libcdep.cc b/lib/sanitizer_common/sanitizer_coverage_libcdep.cc
index 0661e5ff5..efa7db9bd 100644
--- a/lib/sanitizer_common/sanitizer_coverage_libcdep.cc
+++ b/lib/sanitizer_common/sanitizer_coverage_libcdep.cc
@@ -113,6 +113,8 @@ class CoverageData {
   uptr *data();
   uptr size() const;
 
+  void SetPcBuffer(uptr* data, uptr length);
+
  private:
   struct NamedPcRange {
     const char *copied_module_name;
@@ -143,6 +145,9 @@ class CoverageData {
   // Descriptor of the file mapped pc array.
   fd_t pc_fd;
 
+  uptr *pc_buffer;
+  uptr pc_buffer_len;
+
   // Vector of coverage guard arrays, protected by mu.
   InternalMmapVectorNoCtor<s32*> guard_array_vec;
 
@@ -214,6 +219,9 @@ void CoverageData::Enable() {
     atomic_store(&pc_array_size, kPcArrayMaxSize, memory_order_relaxed);
   }
 
+  pc_buffer = nullptr;
+  pc_buffer_len = 0;
+
   cc_array = reinterpret_cast<uptr **>(MmapNoReserveOrDie(
       sizeof(uptr *) * kCcArrayMaxSize, "CovInit::cc_array"));
   atomic_store(&cc_array_size, kCcArrayMaxSize, memory_order_relaxed);
@@ -419,6 +427,7 @@ void CoverageData::Add(uptr pc, u32 *guard) {
            atomic_load(&pc_array_size, memory_order_acquire));
   uptr counter = atomic_fetch_add(&coverage_counter, 1, memory_order_relaxed);
   pc_array[idx] = BundlePcAndCounter(pc, counter);
+  if (pc_buffer && counter < pc_buffer_len) pc_buffer[counter] = pc;
 }
 
 // Registers a pair caller=>callee.
@@ -872,6 +881,11 @@ void CoverageData::DumpAll() {
   DumpCallerCalleePairs();
 }
 
+void CoverageData::SetPcBuffer(uptr* data, uptr length) {
+  pc_buffer = data;
+  pc_buffer_len = length;
+}
+
 void CovPrepareForSandboxing(__sanitizer_sandbox_arguments *args) {
   if (!args) return;
   if (!coverage_enabled) return;
@@ -1007,6 +1021,16 @@ uptr __sanitizer_get_coverage_guards(uptr **data) {
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE
+void __sanitizer_set_coverage_pc_buffer(uptr *data, uptr length) {
+  coverage_data.SetPcBuffer(data, length);
+}
+
+SANITIZER_INTERFACE_ATTRIBUTE
+uptr __sanitizer_get_coverage_pc_buffer_pos() {
+  return __sanitizer_get_total_unique_coverage();
+}
+
+SANITIZER_INTERFACE_ATTRIBUTE
 uptr __sanitizer_get_number_of_counters() {
   return coverage_data.GetNumberOf8bitCounters();
 }
diff --git a/test/asan/TestCases/coverage-pc-buffer.cc b/test/asan/TestCases/coverage-pc-buffer.cc
new file mode 100644
index 000000000..dd9beafb1
--- /dev/null
+++ b/test/asan/TestCases/coverage-pc-buffer.cc
@@ -0,0 +1,49 @@
+// Test __sanitizer_coverage_pc_buffer().
+
+// RUN: %clangxx_asan -fsanitize-coverage=edge %stdcxx11 %s -O3 -o %t && %run %t
+
+// UNSUPPORTED: android
+
+#include <assert.h>
+#include <memory>
+#include <sanitizer/coverage_interface.h>
+#include <stdint.h>
+#include <stdio.h>
+
+static volatile int sink;
+__attribute__((noinline)) void foo() { sink = 1; }
+
+void assertNotZeroPcs(uintptr_t *buf, uintptr_t size) {
+  assert(buf);
+  for (uintptr_t i = 0; i < size; ++i)
+    assert(buf[i]);
+}
+
+int main() {
+  uintptr_t buf_size = 1 << 20;
+  std::unique_ptr<uintptr_t[]> buf(new uintptr_t[buf_size]);
+  __sanitizer_set_coverage_pc_buffer(buf.get(), buf_size);
+
+  {
+    uintptr_t sz = __sanitizer_get_coverage_pc_buffer_pos();
+    assertNotZeroPcs(buf.get(), sz);
+    assert(sz);
+  }
+
+  {
+    uintptr_t sz = __sanitizer_get_coverage_pc_buffer_pos();
+    foo();
+    uintptr_t sz1 = __sanitizer_get_coverage_pc_buffer_pos();
+    assertNotZeroPcs(buf.get(), sz1);
+    assert(sz1 > sz);
+  }
+
+  {
+    uintptr_t sz = __sanitizer_get_coverage_pc_buffer_pos();
+    // reset coverage to 0.
+    __sanitizer_reset_coverage();
+    uintptr_t sz1 = __sanitizer_get_coverage_pc_buffer_pos();
+    assertNotZeroPcs(buf.get(), sz1);
+    assert(sz1 < sz);
+  }
+}