From 5971f18133cbdece6012b3e0f54ae85877e4306f Mon Sep 17 00:00:00 2001
From: Mike Aizatsky <aizatsky@chromium.org>
Date: Fri, 26 Feb 2016 01:17:22 +0000
Subject: [PATCH] [sancov] Pruning full dominator blocks from instrumentation.

Summary:
This is the first simple attempt to reduce number of coverage-
instrumented blocks.

If a basic block dominates all its successors, then its coverage
information is useless to us. Ingore such blocks if
santizer-coverage-prune-tree option is set.

Differential Revision: http://reviews.llvm.org/D17626

llvm-svn: 261949
---
 .../Instrumentation/SanitizerCoverage.cpp     | 36 ++++++++++++++++---
 .../SanitizerCoverage/coverage.ll             |  9 +++++
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp b/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
index 34e8b579c1ca..fd3889e9a740 100644
--- a/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
+++ b/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
@@ -28,13 +28,14 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/Transforms/Instrumentation.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/EHPersonalities.h"
+#include "llvm/IR/CFG.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfo.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InlineAsm.h"
@@ -45,6 +46,7 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Instrumentation.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
@@ -94,6 +96,11 @@ static cl::opt<bool>
                                       "instructions"),
                              cl::Hidden, cl::init(false));
 
+static cl::opt<bool> ClPruneBlocks(
+    "sanitizer-coverage-prune-blocks",
+    cl::desc("Reduce the number of instrumented blocks (experimental)"),
+    cl::Hidden, cl::init(false));
+
 // Experimental 8-bit counters used as an additional search heuristic during
 // coverage-guided fuzzing.
 // The counters are not thread-friendly:
@@ -298,6 +305,22 @@ bool SanitizerCoverageModule::runOnModule(Module &M) {
   return true;
 }
 
+static bool shouldInstrumentBlock(const BasicBlock *BB,
+                                  const DominatorTree *DT) {
+  if (!ClPruneBlocks)
+    return true;
+  if (succ_begin(BB) == succ_end(BB))
+    return true;
+
+  // Check if BB dominates all its successors.
+  for (const BasicBlock *SUCC : make_range(succ_begin(BB), succ_end(BB))) {
+    if (!DT->dominates(BB, SUCC))
+      return true;
+  }
+
+  return false;
+}
+
 bool SanitizerCoverageModule::runOnFunction(Function &F) {
   if (F.empty()) return false;
   if (F.getName().find(".module_ctor") != std::string::npos)
@@ -311,11 +334,15 @@ bool SanitizerCoverageModule::runOnFunction(Function &F) {
   if (Options.CoverageType >= SanitizerCoverageOptions::SCK_Edge)
     SplitAllCriticalEdges(F);
   SmallVector<Instruction*, 8> IndirCalls;
-  SmallVector<BasicBlock*, 16> AllBlocks;
+  SmallVector<BasicBlock *, 16> BlocksToInstrument;
   SmallVector<Instruction*, 8> CmpTraceTargets;
   SmallVector<Instruction*, 8> SwitchTraceTargets;
+
+  DominatorTree DT;
+  DT.recalculate(F);
   for (auto &BB : F) {
-    AllBlocks.push_back(&BB);
+    if (shouldInstrumentBlock(&BB, &DT))
+      BlocksToInstrument.push_back(&BB);
     for (auto &Inst : BB) {
       if (Options.IndirectCalls) {
         CallSite CS(&Inst);
@@ -330,7 +357,8 @@ bool SanitizerCoverageModule::runOnFunction(Function &F) {
       }
     }
   }
-  InjectCoverage(F, AllBlocks);
+
+  InjectCoverage(F, BlocksToInstrument);
   InjectCoverageForIndirectCalls(F, IndirCalls);
   InjectTraceForCmp(F, CmpTraceTargets);
   InjectTraceForSwitch(F, SwitchTraceTargets);
diff --git a/llvm/test/Instrumentation/SanitizerCoverage/coverage.ll b/llvm/test/Instrumentation/SanitizerCoverage/coverage.ll
index f7fa983a0021..71177eb33a1a 100644
--- a/llvm/test/Instrumentation/SanitizerCoverage/coverage.ll
+++ b/llvm/test/Instrumentation/SanitizerCoverage/coverage.ll
@@ -13,6 +13,7 @@
 ; RUN:      -S | FileCheck %s --check-prefix=CHECK2
 ; RUN: opt < %s -sancov -sanitizer-coverage-level=2 -sanitizer-coverage-block-threshold=1 \
 ; RUN:      -S | FileCheck %s --check-prefix=CHECK_WITH_CHECK
+; RUN: opt < %s -sancov -sanitizer-coverage-level=2 -sanitizer-coverage-prune-blocks=1 -S | FileCheck %s --check-prefix=CHECKPRUNE
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64"
 target triple = "x86_64-unknown-linux-gnu"
@@ -134,3 +135,11 @@ entry:
 ; CHECK4-LABEL: define void @call_unreachable
 ; CHECK4-NOT: __sanitizer_cov
 ; CHECK4: unreachable
+
+; CHECKPRUNE-LABEL: define void @foo
+; CHECKPRUNE: call void @__sanitizer_cov
+; CHECKPRUNE: call void asm sideeffect "", ""()
+; CHECKPRUNE: call void @__sanitizer_cov
+; CHECKPRUNE: call void asm sideeffect "", ""()
+; CHECKPRUNE-NOT: call void @__sanitizer_cov
+; CHECKPRUNE: ret void