1 files changed, 19 insertions, 7 deletions

diff --git a/llvm/lib/Transforms/Scalar/BoundsChecking.cpp b/llvm/lib/Transforms/Scalar/BoundsChecking.cpp
index 004f34f3bd6..85c5e111e0f 100644
--- a/llvm/lib/Transforms/Scalar/BoundsChecking.cpp
+++ b/llvm/lib/Transforms/Scalar/BoundsChecking.cpp
@@ -62,6 +62,7 @@ namespace {
     unsigned Penalty;
 
     BasicBlock *getTrapBB();
+    void emitBranchToTrap(Value *Cmp = 0);
     ConstTriState computeAllocSize(Value *Alloc, uint64_t &Size,
                                    Value* &SizeValue);
     bool instrument(Value *Ptr, Value *Val);
@@ -94,6 +95,22 @@ BasicBlock *BoundsChecking::getTrapBB() {
 }
 
 
+/// emitBranchToTrap - emit a branch instruction to a trap block.
+/// If Cmp is non-null, perform a jump only if its value evaluates to true.
+void BoundsChecking::emitBranchToTrap(Value *Cmp) {
+  Instruction *Inst = Builder->GetInsertPoint();
+  BasicBlock *OldBB = Inst->getParent();
+  BasicBlock *Cont = OldBB->splitBasicBlock(Inst);
+  OldBB->getTerminator()->eraseFromParent();
+
+  // FIXME: add unlikely branch taken metadata?
+  if (Cmp)
+    BranchInst::Create(getTrapBB(), Cont, Cmp, OldBB);
+  else
+    BranchInst::Create(getTrapBB(), OldBB);
+}
+
+
 /// computeAllocSize - compute the object size allocated by an allocation
 /// site. Returns NotConst if the size is not constant (in SizeValue), Const if
 /// the size is constant (in Size), and Dunno if the size could not be
@@ -254,7 +271,7 @@ bool BoundsChecking::instrument(Value *Ptr, Value *InstVal) {
   if (!OffsetValue && ConstAlloc == Const) {
     if (Size < Offset || (Size - Offset) < NeededSize) {
       // Out of bounds
-      Builder->CreateBr(getTrapBB());
+      emitBranchToTrap();
       ++ChecksAdded;
       return true;
     }
@@ -278,13 +295,8 @@ bool BoundsChecking::instrument(Value *Ptr, Value *InstVal) {
   Value *Cmp1 = Builder->CreateICmpULT(SizeValue, OffsetValue);
   Value *Cmp2 = Builder->CreateICmpULT(ObjSize, NeededSizeVal);
   Value *Or = Builder->CreateOr(Cmp1, Cmp2);
+  emitBranchToTrap(Or);
 
-  // FIXME: add unlikely branch taken metadata?
-  Instruction *Inst = Builder->GetInsertPoint();
-  BasicBlock *OldBB = Inst->getParent();
-  BasicBlock *Cont = OldBB->splitBasicBlock(Inst);
-  OldBB->getTerminator()->eraseFromParent();
-  BranchInst::Create(getTrapBB(), Cont, Or, OldBB);
   ++ChecksAdded;
   return true;
 }