2 files changed, 35 insertions, 22 deletions

diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
index 14348b9bac6..336802668ca 100644
--- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
@@ -18,6 +18,7 @@
 #include "FunctionBlackList.h"
 #include "llvm/Function.h"
 #include "llvm/IRBuilder.h"
+#include "llvm/InlineAsm.h"
 #include "llvm/IntrinsicInst.h"
 #include "llvm/LLVMContext.h"
 #include "llvm/Module.h"
@@ -224,6 +225,7 @@ struct AddressSanitizer : public ModulePass {
   OwningPtr<FunctionBlackList> BL;
   // This array is indexed by AccessIsWrite and log2(AccessSize).
   Function *AsanErrorCallback[2][kNumberOfAccessSizes];
+  InlineAsm *EmptyAsm;
 };
 
 }  // namespace
@@ -276,7 +278,7 @@ static BranchInst *splitBlockAndInsertIfThen(Value *Cmp,
   BranchInst *CheckTerm = 0;
   if (!ThenBlock) {
     LLVMContext &C = Head->getParent()->getParent()->getContext();
-    ThenBlock = BasicBlock::Create(C, "", Head->getParent());
+    ThenBlock = BasicBlock::Create(C, "", Head->getParent(), Tail);
     CheckTerm = BranchInst::Create(Tail, ThenBlock);
   }
   BranchInst *HeadNewTerm =
@@ -414,7 +416,10 @@ Instruction *AddressSanitizer::generateCrashCode(
                            Addr, PC);
   else
     Call = IRB.CreateCall(AsanErrorCallback[IsWrite][AccessSizeIndex], Addr);
-  Call->setDoesNotReturn();
+  // We don't do Call->setDoesNotReturn() because the BB already has
+  // UnreachableInst at the end.
+  // This EmptyAsm is required to avoid callback merge.
+  IRB.CreateCall(EmptyAsm);
   return Call;
 }
 
@@ -483,10 +488,13 @@ void AddressSanitizer::instrumentAddress(AsanFunctionContext &AFC,
 
   size_t Granularity = 1 << MappingScale;
   if (TypeSize < 8 * Granularity) {
-    Instruction *CheckTerm = splitBlockAndInsertIfThen(Cmp);
+    BranchInst *CheckTerm = splitBlockAndInsertIfThen(Cmp);
+    assert(CheckTerm->isUnconditional());
+    BasicBlock *NextBB = CheckTerm->getSuccessor(0);
     IRB.SetInsertPoint(CheckTerm);
     Value *Cmp2 = createSlowPathCmp(IRB, AddrLong, ShadowValue, TypeSize);
-    splitBlockAndInsertIfThen(Cmp2, CrashBlock);
+    BranchInst *NewTerm = BranchInst::Create(CrashBlock, NextBB, Cmp2);
+    ReplaceInstWithInst(CheckTerm, NewTerm);
   } else {
     splitBlockAndInsertIfThen(Cmp, CrashBlock);
   }
@@ -695,6 +703,10 @@ bool AddressSanitizer::runOnModule(Module &M) {
           M.getOrInsertFunction(FunctionName, IRB.getVoidTy(), IntptrTy, NULL));
     }
   }
+  // We insert an empty inline asm after __asan_report* to avoid callback merge.
+  EmptyAsm = InlineAsm::get(FunctionType::get(IRB.getVoidTy(), false),
+                            StringRef(""), StringRef(""),
+                            /*hasSideEffects=*/true);
 
   llvm::Triple targetTriple(M.getTargetTriple());
   bool isAndroid = targetTriple.getEnvironment() == llvm::Triple::ANDROIDEABI;
diff --git a/llvm/test/Instrumentation/AddressSanitizer/basic.ll b/llvm/test/Instrumentation/AddressSanitizer/basic.ll
index 183cddcb5ca..294ca8ab756 100644
--- a/llvm/test/Instrumentation/AddressSanitizer/basic.ll
+++ b/llvm/test/Instrumentation/AddressSanitizer/basic.ll
@@ -16,15 +16,6 @@ define i32 @test_load(i32* %a) address_safety {
 ; CHECK:   icmp ne i8
 ; CHECK:   br i1 %{{.*}}, label %{{.*}}, label %{{.*}}
 ;
-; The actual load comes next because ASan adds the last instrumentation block
-; to the end of the function.
-; CHECK:   %tmp1 = load i32* %a
-; CHECK:   ret i32 %tmp1
-
-; The crash block reports the error.
-; CHECK:   call void @__asan_report_load4(i64 %[[LOAD_ADDR]]) noreturn
-; CHECK:   unreachable
-;
 ; First instrumentation block refines the shadow test.
 ; CHECK:   and i64 %[[LOAD_ADDR]], 7
 ; CHECK:   add i64 %{{.*}}, 3
@@ -32,6 +23,16 @@ define i32 @test_load(i32* %a) address_safety {
 ; CHECK:   icmp sge i8 %{{.*}}, %[[LOAD_SHADOW]]
 ; CHECK:   br i1 %{{.*}}, label %{{.*}}, label %{{.*}}
 ;
+; The actual load comes next because ASan adds the crash block
+; to the end of the function.
+; CHECK:   %tmp1 = load i32* %a
+; CHECK:   ret i32 %tmp1
+
+; The crash block reports the error.
+; CHECK:   call void @__asan_report_load4(i64 %[[LOAD_ADDR]])
+; CHECK:   unreachable
+;
+
 
 entry:
   %tmp1 = load i32* %a
@@ -49,15 +50,6 @@ define void @test_store(i32* %a) address_safety {
 ; CHECK:   icmp ne i8
 ; CHECK:   br i1 %{{.*}}, label %{{.*}}, label %{{.*}}
 ;
-; The actual store comes next because ASan adds the last instrumentation block
-; to the end of the function.
-; CHECK:   store i32 42, i32* %a
-; CHECK:   ret void
-;
-; The crash block reports the error.
-; CHECK:   call void @__asan_report_store4(i64 %[[STORE_ADDR]]) noreturn
-; CHECK:   unreachable
-;
 ; First instrumentation block refines the shadow test.
 ; CHECK:   and i64 %[[STORE_ADDR]], 7
 ; CHECK:   add i64 %{{.*}}, 3
@@ -65,6 +57,15 @@ define void @test_store(i32* %a) address_safety {
 ; CHECK:   icmp sge i8 %{{.*}}, %[[STORE_SHADOW]]
 ; CHECK:   br i1 %{{.*}}, label %{{.*}}, label %{{.*}}
 ;
+; The actual load comes next because ASan adds the crash block
+; to the end of the function.
+; CHECK:   store i32 42, i32* %a
+; CHECK:   ret void
+;
+; The crash block reports the error.
+; CHECK:   call void @__asan_report_store4(i64 %[[STORE_ADDR]])
+; CHECK:   unreachable
+;
 
 entry:
   store i32 42, i32* %a