[sanitizer] fix instrumentation with -mllvm -sanitizer-coverage-block-threshold=0 to actually do something useful.

llvm-svn: 231736

4 files changed, 19 insertions, 8 deletions

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_coverage_libcdep.cc b/compiler-rt/lib/sanitizer_common/sanitizer_coverage_libcdep.cc
index 7b7225f187c..112e4896cf7 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_coverage_libcdep.cc
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_coverage_libcdep.cc
@@ -812,7 +812,8 @@ SANITIZER_INTERFACE_ATTRIBUTE void __sanitizer_cov(u32 *guard) {
 }
 SANITIZER_INTERFACE_ATTRIBUTE void __sanitizer_cov_with_check(u32 *guard) {
   atomic_uint32_t *atomic_guard = reinterpret_cast<atomic_uint32_t*>(guard);
-  if (__sanitizer::atomic_load(atomic_guard, memory_order_relaxed))
+  if (static_cast<s32>(
+          __sanitizer::atomic_load(atomic_guard, memory_order_relaxed)) < 0)
     __sanitizer_cov(guard);
 }
 SANITIZER_INTERFACE_ATTRIBUTE void
diff --git a/compiler-rt/test/asan/TestCases/Linux/coverage-levels.cc b/compiler-rt/test/asan/TestCases/Linux/coverage-levels.cc
index f84fd0bf6a9..63bed650ba0 100644
--- a/compiler-rt/test/asan/TestCases/Linux/coverage-levels.cc
+++ b/compiler-rt/test/asan/TestCases/Linux/coverage-levels.cc
@@ -6,6 +6,8 @@
 // RUN: ASAN_OPTIONS=coverage=1:coverage_bitset=1:verbosity=1 %run %t 2>&1 | FileCheck %s --check-prefix=CHECK2
 // RUN: %clangxx_asan -O1 -fsanitize-coverage=3  %s -o %t
 // RUN: ASAN_OPTIONS=coverage=1:coverage_bitset=1:verbosity=1 %run %t 2>&1 | FileCheck %s --check-prefix=CHECK3
+// RUN: %clangxx_asan -O1 -fsanitize-coverage=3 -mllvm -sanitizer-coverage-block-threshold=0 %s -o %t
+// RUN: ASAN_OPTIONS=coverage=1:coverage_bitset=1:verbosity=1 %run %t 2>&1 | FileCheck %s --check-prefix=CHECK3
 // RUN: %clangxx_asan -O1 -fsanitize-coverage=3 -mllvm -sanitizer-coverage-8bit-counters=1 %s -o %t
 // RUN: ASAN_OPTIONS=coverage=1:coverage_counters=1:verbosity=1 %run %t 2>&1 | FileCheck %s --check-prefix=CHECK_COUNTERS
 
diff --git a/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp b/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
index fcb6ab04f5e..9c74b5ea5c2 100644
--- a/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
+++ b/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
@@ -111,6 +111,9 @@ class SanitizerCoverageModule : public ModulePass {
                       ArrayRef<Instruction *> IndirCalls);
   void SetNoSanitizeMetada(Instruction *I);
   void InjectCoverageAtBlock(Function &F, BasicBlock &BB, bool UseCalls);
+  unsigned NumberOfInstrumentedBlocks() {
+    return SanCovFunction->getNumUses() + SanCovWithCheckFunction->getNumUses();
+  }
   Function *SanCovFunction;
   Function *SanCovWithCheckFunction;
   Function *SanCovIndirCallFunction;
@@ -192,10 +195,11 @@ bool SanitizerCoverageModule::runOnModule(Module &M) {
   for (auto &F : M)
     runOnFunction(F);
 
+  auto N = NumberOfInstrumentedBlocks();
+
   // Now we know how many elements we need. Create an array of guards
   // with one extra element at the beginning for the size.
-  Type *Int32ArrayNTy =
-      ArrayType::get(Int32Ty, SanCovFunction->getNumUses() + 1);
+  Type *Int32ArrayNTy = ArrayType::get(Int32Ty, N + 1);
   GlobalVariable *RealGuardArray = new GlobalVariable(
       M, Int32ArrayNTy, false, GlobalValue::PrivateLinkage,
       Constant::getNullValue(Int32ArrayNTy), "__sancov_gen_cov");
@@ -211,8 +215,7 @@ bool SanitizerCoverageModule::runOnModule(Module &M) {
     // Make sure the array is 16-aligned.
     static const int kCounterAlignment = 16;
     Type *Int8ArrayNTy =
-        ArrayType::get(Int8Ty, RoundUpToAlignment(SanCovFunction->getNumUses(),
-                                                  kCounterAlignment));
+        ArrayType::get(Int8Ty, RoundUpToAlignment(N, kCounterAlignment));
     RealEightBitCounterArray = new GlobalVariable(
         M, Int8ArrayNTy, false, GlobalValue::PrivateLinkage,
         Constant::getNullValue(Int8ArrayNTy), "__sancov_gen_cov_counter");
@@ -233,7 +236,7 @@ bool SanitizerCoverageModule::runOnModule(Module &M) {
   IRB.SetInsertPoint(CtorFunc->getEntryBlock().getTerminator());
   IRB.CreateCall4(
       SanCovModuleInit, IRB.CreatePointerCast(RealGuardArray, Int32PtrTy),
-      ConstantInt::get(IntptrTy, SanCovFunction->getNumUses()),
+      ConstantInt::get(IntptrTy, N),
       ClUse8bitCounters
           ? IRB.CreatePointerCast(RealEightBitCounterArray, Int8PtrTy)
           : Constant::getNullValue(Int8PtrTy),
@@ -333,7 +336,7 @@ void SanitizerCoverageModule::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
   SmallVector<Value *, 1> Indices;
   Value *GuardP = IRB.CreateAdd(
       IRB.CreatePointerCast(GuardArray, IntptrTy),
-      ConstantInt::get(IntptrTy, (1 + SanCovFunction->getNumUses()) * 4));
+      ConstantInt::get(IntptrTy, (1 + NumberOfInstrumentedBlocks()) * 4));
   Type *Int32PtrTy = PointerType::getUnqual(IRB.getInt32Ty());
   GuardP = IRB.CreateIntToPtr(GuardP, Int32PtrTy);
   if (UseCalls) {
@@ -357,7 +360,7 @@ void SanitizerCoverageModule::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
     IRB.SetInsertPoint(IP);
     Value *P = IRB.CreateAdd(
         IRB.CreatePointerCast(EightBitCounterArray, IntptrTy),
-        ConstantInt::get(IntptrTy, SanCovFunction->getNumUses() - 1));
+        ConstantInt::get(IntptrTy, NumberOfInstrumentedBlocks() - 1));
     P = IRB.CreateIntToPtr(P, IRB.getInt8PtrTy());
     LoadInst *LI = IRB.CreateLoad(P);
     Value *Inc = IRB.CreateAdd(LI, ConstantInt::get(IRB.getInt8Ty(), 1));
diff --git a/llvm/test/Instrumentation/SanitizerCoverage/coverage.ll b/llvm/test/Instrumentation/SanitizerCoverage/coverage.ll
index 024a9dceff3..b2f0ab0680b 100644
--- a/llvm/test/Instrumentation/SanitizerCoverage/coverage.ll
+++ b/llvm/test/Instrumentation/SanitizerCoverage/coverage.ll
@@ -2,6 +2,7 @@
 ; RUN: opt < %s -sancov -sanitizer-coverage-level=1 -S | FileCheck %s --check-prefix=CHECK1
 ; RUN: opt < %s -sancov -sanitizer-coverage-level=2 -S | FileCheck %s --check-prefix=CHECK2
 ; RUN: opt < %s -sancov -sanitizer-coverage-level=2 -sanitizer-coverage-block-threshold=10 -S | FileCheck %s --check-prefix=CHECK2
+; RUN: opt < %s -sancov -sanitizer-coverage-level=2 -sanitizer-coverage-block-threshold=0  -S | FileCheck %s --check-prefix=CHECK_WITH_CHECK
 ; RUN: opt < %s -sancov -sanitizer-coverage-level=2 -sanitizer-coverage-block-threshold=1  -S | FileCheck %s --check-prefix=CHECK_WITH_CHECK
 ; RUN: opt < %s -sancov -sanitizer-coverage-level=3 -sanitizer-coverage-block-threshold=10 -S | FileCheck %s --check-prefix=CHECK3
 ; RUN: opt < %s -sancov -sanitizer-coverage-level=4 -S | FileCheck %s --check-prefix=CHECK4
@@ -51,6 +52,10 @@ entry:
 ; CHECK_WITH_CHECK-LABEL: define void @foo
 ; CHECK_WITH_CHECK: __sanitizer_cov_with_check
 ; CHECK_WITH_CHECK: ret void
+; CHECK_WITH_CHECK-LABEL: define internal void @sancov.module_ctor
+; CHECK_WITH_CHECK-NOT: ret
+; CHECK_WITH_CHECK: call void @__sanitizer_cov_module_init({{.*}}, i64 4,
+; CHECK_WITH_CHECK: ret
 
 ; CHECK2-LABEL: define void @foo
 ; CHECK2: call void @__sanitizer_cov