From 73522d167890e904e6a42eb06836dd1ecc2fb8fc Mon Sep 17 00:00:00 2001
From: Vedant Kumar <vsk@apple.com>
Date: Tue, 26 Feb 2019 22:55:46 +0000
Subject: [PATCH] [HotColdSplit] Disable splitting for sanitized functions

Splitting can make sanitizer errors harder to understand, as the
trapping instruction may not be in the function where the bug was
detected.

rdar://48142697

llvm-svn: 354931
---
 llvm/lib/Transforms/IPO/HotColdSplitting.cpp       | 11 +++-
 .../Transforms/HotColdSplit/X86/do-not-split.ll    | 72 ++++++++++++++++++++++
 2 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/HotColdSplitting.cpp b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
index b8def7a..a237eb7 100644
--- a/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
+++ b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
@@ -111,10 +111,11 @@ bool unlikelyExecuted(BasicBlock &BB) {
   if (BB.isEHPad() || isa<ResumeInst>(BB.getTerminator()))
     return true;
 
-  // The block is cold if it calls/invokes a cold function.
+  // The block is cold if it calls/invokes a cold function. However, do not
+  // mark sanitizer traps as cold.
   for (Instruction &I : BB)
     if (auto CS = CallSite(&I))
-      if (CS.hasFnAttr(Attribute::Cold))
+      if (CS.hasFnAttr(Attribute::Cold) && !CS->getMetadata("nosanitize"))
         return true;
 
   // The block is cold if it has an unreachable terminator, unless it's
@@ -235,6 +236,12 @@ bool HotColdSplitting::shouldOutlineFrom(const Function &F) const {
   if (F.hasFnAttribute(Attribute::NoInline))
     return false;
 
+  if (F.hasFnAttribute(Attribute::SanitizeAddress) ||
+      F.hasFnAttribute(Attribute::SanitizeHWAddress) ||
+      F.hasFnAttribute(Attribute::SanitizeThread) ||
+      F.hasFnAttribute(Attribute::SanitizeMemory))
+    return false;
+
   return true;
 }
 
diff --git a/llvm/test/Transforms/HotColdSplit/X86/do-not-split.ll b/llvm/test/Transforms/HotColdSplit/X86/do-not-split.ll
index 448e63a..076174d 100644
--- a/llvm/test/Transforms/HotColdSplit/X86/do-not-split.ll
+++ b/llvm/test/Transforms/HotColdSplit/X86/do-not-split.ll
@@ -110,6 +110,78 @@ if.end:                                           ; preds = %entry
   ret void
 }
 
+; CHECK-LABEL: @sanitize_address
+; CHECK-NOT: sanitize_address.cold.1
+define void @sanitize_address() sanitize_address {
+entry:
+  br i1 undef, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  call void @sink()
+  ret void
+
+if.end:                                           ; preds = %entry
+  ret void
+}
+
+; CHECK-LABEL: @sanitize_hwaddress
+; CHECK-NOT: sanitize_hwaddress.cold.1
+define void @sanitize_hwaddress() sanitize_hwaddress {
+entry:
+  br i1 undef, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  call void @sink()
+  ret void
+
+if.end:                                           ; preds = %entry
+  ret void
+}
+
+; CHECK-LABEL: @sanitize_thread
+; CHECK-NOT: sanitize_thread.cold.1
+define void @sanitize_thread() sanitize_thread {
+entry:
+  br i1 undef, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  call void @sink()
+  ret void
+
+if.end:                                           ; preds = %entry
+  ret void
+}
+
+; CHECK-LABEL: @sanitize_memory
+; CHECK-NOT: sanitize_memory.cold.1
+define void @sanitize_memory() sanitize_memory {
+entry:
+  br i1 undef, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  call void @sink()
+  ret void
+
+if.end:                                           ; preds = %entry
+  ret void
+}
+
+declare void @llvm.trap() cold noreturn
+
+; CHECK-LABEL: @nosanitize_call
+; CHECK-NOT: nosanitize_call.cold.1
+define void @nosanitize_call() sanitize_memory {
+entry:
+  br i1 undef, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  call void @llvm.trap(), !nosanitize !2
+  unreachable
+
+if.end:                                           ; preds = %entry
+  ret void
+}
+
 declare void @llvm.dbg.value(metadata, metadata, metadata)
 
 declare void @sink() cold
-- 
2.7.4