7 files changed, 46 insertions, 5 deletions
diff --git a/compiler-rt/lib/esan/cache_frag.cpp b/compiler-rt/lib/esan/cache_frag.cpp
index 343abfcd584..a3e612daceb 100644
--- a/compiler-rt/lib/esan/cache_frag.cpp
+++ b/compiler-rt/lib/esan/cache_frag.cpp
@@ -199,4 +199,10 @@ int finalizeCacheFrag() {
   return 0;
 }
 
+void reportCacheFrag() {
+  VPrintf(2, "in esan::%s\n", __FUNCTION__);
+  // FIXME: Not yet implemented.  We need to iterate over all of the
+  // compilation unit data.
+}
+
 } // namespace __esan
diff --git a/compiler-rt/lib/esan/cache_frag.h b/compiler-rt/lib/esan/cache_frag.h
index 9f879322074..646d3f85ed9 100644
--- a/compiler-rt/lib/esan/cache_frag.h
+++ b/compiler-rt/lib/esan/cache_frag.h
@@ -22,6 +22,7 @@ void processCacheFragCompilationUnitExit(void *Ptr);
 
 void initializeCacheFrag();
 int finalizeCacheFrag();
+void reportCacheFrag();
 
 } // namespace __esan
 
diff --git a/compiler-rt/lib/esan/esan.cpp b/compiler-rt/lib/esan/esan.cpp
index 970f91a837d..3c69b4e91b5 100644
--- a/compiler-rt/lib/esan/esan.cpp
+++ b/compiler-rt/lib/esan/esan.cpp
@@ -228,6 +228,15 @@ int finalizeLibrary() {
   return 0;
 }
 
+void reportResults() {
+  VPrintf(1, "in esan::%s\n", __FUNCTION__);
+  if (__esan_which_tool == ESAN_CacheFrag) {
+    return reportCacheFrag();
+  } else if (__esan_which_tool == ESAN_WorkingSet) {
+    return reportWorkingSet();
+  }
+}
+
 void processCompilationUnitInit(void *Ptr) {
   VPrintf(2, "in esan::%s\n", __FUNCTION__);
   if (__esan_which_tool == ESAN_CacheFrag) {
diff --git a/compiler-rt/lib/esan/esan.h b/compiler-rt/lib/esan/esan.h
index 4171522dbbf..371810d5eba 100644
--- a/compiler-rt/lib/esan/esan.h
+++ b/compiler-rt/lib/esan/esan.h
@@ -37,6 +37,7 @@ extern bool EsanDuringInit;
 
 void initializeLibrary(ToolType Tool);
 int finalizeLibrary();
+void reportResults();
 // Esan creates the variable per tool per compilation unit at compile time
 // and passes its pointer Ptr to the runtime library.
 void processCompilationUnitInit(void *Ptr);
diff --git a/compiler-rt/lib/esan/esan_interface.cpp b/compiler-rt/lib/esan/esan_interface.cpp
index f6ad3caa841..8a64d152611 100644
--- a/compiler-rt/lib/esan/esan_interface.cpp
+++ b/compiler-rt/lib/esan/esan_interface.cpp
@@ -109,3 +109,10 @@ void __esan_unaligned_loadN(void *Addr, uptr Size) {
 void __esan_unaligned_storeN(void *Addr, uptr Size) {
   processRangeAccess(GET_CALLER_PC(), (uptr)Addr, Size, true);
 }
+
+// Public interface:
+extern "C" {
+SANITIZER_INTERFACE_ATTRIBUTE void __esan_report() {
+  reportResults();
+}
+} // extern "C"
diff --git a/compiler-rt/lib/esan/working_set.cpp b/compiler-rt/lib/esan/working_set.cpp
index 622fd29e846..3fde5a8b5fc 100644
--- a/compiler-rt/lib/esan/working_set.cpp
+++ b/compiler-rt/lib/esan/working_set.cpp
@@ -118,6 +118,8 @@ void processRangeAccessWorkingSet(uptr PC, uptr Addr, SIZE_T Size,
 }
 
 // This routine will word-align ShadowStart and ShadowEnd prior to scanning.
+// It does *not* clear for BitIdx==TotalWorkingSetBitIdx, as that top bit
+// measures the access during the entire execution and should never be cleared.
 static u32 countAndClearShadowValues(u32 BitIdx, uptr ShadowStart,
                                      uptr ShadowEnd) {
   u32 WorkingSetSize = 0;
@@ -127,6 +129,8 @@ static u32 countAndClearShadowValues(u32 BitIdx, uptr ShadowStart,
   // Get word aligned start.
   ShadowStart = RoundDownTo(ShadowStart, sizeof(u32));
   bool Accum = getFlags()->record_snapshots && BitIdx < MaxAccumBitIdx;
+  // Do not clear the bit that measures access during the entire execution.
+  bool Clear = BitIdx < TotalWorkingSetBitIdx;
   for (u32 *Ptr = (u32 *)ShadowStart; Ptr < (u32 *)ShadowEnd; ++Ptr) {
     if ((*Ptr & WordValue) != 0) {
       byte *BytePtr = (byte *)Ptr;
@@ -139,8 +143,10 @@ static u32 countAndClearShadowValues(u32 BitIdx, uptr ShadowStart,
           }
         }
       }
-      // Clear this bit from every shadow byte.
-      *Ptr &= ~WordValue;
+      if (Clear) {
+        // Clear this bit from every shadow byte.
+        *Ptr &= ~WordValue;
+      }
     }
   }
   return WorkingSetSize;
@@ -149,6 +155,8 @@ static u32 countAndClearShadowValues(u32 BitIdx, uptr ShadowStart,
 // Scan shadow memory to calculate the number of cache lines being accessed,
 // i.e., the number of non-zero bits indexed by BitIdx in each shadow byte.
 // We also clear the lowest bits (most recent working set snapshot).
+// We do *not* clear for BitIdx==TotalWorkingSetBitIdx, as that top bit
+// measures the access during the entire execution and should never be cleared.
 static u32 computeWorkingSizeAndReset(u32 BitIdx) {
   u32 WorkingSetSize = 0;
   MemoryMappingLayout MemIter(true/*cache*/);
@@ -226,10 +234,9 @@ static u32 getSizeForPrinting(u32 NumOfCachelines, const char *&Unit) {
   }
 }
 
-int finalizeWorkingSet() {
+void reportWorkingSet() {
   const char *Unit;
   if (getFlags()->record_snapshots) {
-    Thread.joinThread();
     u32 Freq = 1;
     Report(" Total number of samples: %u\n", SnapshotNum);
     for (u32 i = 0; i < NumFreq; ++i) {
@@ -243,7 +250,6 @@ int finalizeWorkingSet() {
                SizePerFreq[i][j]);
       }
       Freq = Freq << getFlags()->snapshot_step;
-      SizePerFreq[i].free();
     }
   }
 
@@ -252,6 +258,16 @@ int finalizeWorkingSet() {
   u32 Size = getSizeForPrinting(NumOfCachelines, Unit);
   Report(" %s: the total working set size: %u %s (%u cache lines)\n",
          SanitizerToolName, Size, Unit, NumOfCachelines);
+}
+
+int finalizeWorkingSet() {
+  if (getFlags()->record_snapshots)
+    Thread.joinThread();
+  reportWorkingSet();
+  if (getFlags()->record_snapshots) {
+    for (u32 i = 0; i < NumFreq; ++i)
+      SizePerFreq[i].free();
+  }
   return 0;
 }
 
diff --git a/compiler-rt/lib/esan/working_set.h b/compiler-rt/lib/esan/working_set.h
index 660b5d066a7..38ff0635d0a 100644
--- a/compiler-rt/lib/esan/working_set.h
+++ b/compiler-rt/lib/esan/working_set.h
@@ -23,6 +23,7 @@ namespace __esan {
 void initializeWorkingSet();
 void initializeShadowWorkingSet();
 int finalizeWorkingSet();
+void reportWorkingSet();
 void processRangeAccessWorkingSet(uptr PC, uptr Addr, SIZE_T Size,
                                   bool IsWrite);