1 files changed, 27 insertions, 16 deletions

diff --git a/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp b/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp
index 9898c131c1f..741f07ecb4a 100644
--- a/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp
+++ b/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp
@@ -91,8 +91,9 @@ namespace {
                         SmallPtrSetImpl<Instruction *> &DependingInstructions,
                         SmallPtrSetImpl<const BasicBlock *> &Visited);
 
-    void tryToContractReleaseIntoStoreStrong(Instruction *Release,
-                                             inst_iterator &Iter);
+    void tryToContractReleaseIntoStoreStrong(
+        Instruction *Release, inst_iterator &Iter,
+        const DenseMap<BasicBlock *, ColorVector> &BlockColors);
 
     void getAnalysisUsage(AnalysisUsage &AU) const override;
     bool doInitialization(Module &M) override;
@@ -306,6 +307,24 @@ findRetainForStoreStrongContraction(Value *New, StoreInst *Store,
   return Retain;
 }
 
+/// Create a call instruction with the correct funclet token. Should be used
+/// instead of calling CallInst::Create directly.
+static CallInst *
+createCallInst(Value *Func, ArrayRef<Value *> Args, const Twine &NameStr,
+               Instruction *InsertBefore,
+               const DenseMap<BasicBlock *, ColorVector> &BlockColors) {
+  SmallVector<OperandBundleDef, 1> OpBundles;
+  if (!BlockColors.empty()) {
+    const ColorVector &CV = BlockColors.find(InsertBefore->getParent())->second;
+    assert(CV.size() == 1 && "non-unique color for block!");
+    Instruction *EHPad = CV.front()->getFirstNonPHI();
+    if (EHPad->isEHPad())
+      OpBundles.emplace_back("funclet", EHPad);
+  }
+
+  return CallInst::Create(Func, Args, OpBundles, NameStr, InsertBefore);
+}
+
 /// Attempt to merge an objc_release with a store, load, and objc_retain to form
 /// an objc_storeStrong. An objc_storeStrong:
 ///
@@ -333,8 +352,9 @@ findRetainForStoreStrongContraction(Value *New, StoreInst *Store,
 ///     (4).
 ///  2. We need to make sure that any re-orderings of (1), (2), (3), (4) are
 ///     safe.
-void ObjCARCContract::tryToContractReleaseIntoStoreStrong(Instruction *Release,
-                                                          inst_iterator &Iter) {
+void ObjCARCContract::tryToContractReleaseIntoStoreStrong(
+    Instruction *Release, inst_iterator &Iter,
+    const DenseMap<BasicBlock *, ColorVector> &BlockColors) {
   // See if we are releasing something that we just loaded.
   auto *Load = dyn_cast<LoadInst>(GetArgRCIdentityRoot(Release));
   if (!Load || !Load->isSimple())
@@ -386,7 +406,7 @@ void ObjCARCContract::tryToContractReleaseIntoStoreStrong(Instruction *Release,
   if (Args[1]->getType() != I8X)
     Args[1] = new BitCastInst(Args[1], I8X, "", Store);
   Constant *Decl = EP.get(ARCRuntimeEntryPointKind::StoreStrong);
-  CallInst *StoreStrong = CallInst::Create(Decl, Args, "", Store);
+  CallInst *StoreStrong = createCallInst(Decl, Args, "", Store, BlockColors);
   StoreStrong->setDoesNotThrow();
   StoreStrong->setDebugLoc(Store->getDebugLoc());
 
@@ -462,16 +482,7 @@ bool ObjCARCContract::tryToPeepholeInstruction(
             RVInstMarker->getString(),
             /*Constraints=*/"", /*hasSideEffects=*/true);
 
-        SmallVector<OperandBundleDef, 1> OpBundles;
-        if (!BlockColors.empty()) {
-          const ColorVector &CV = BlockColors.find(Inst->getParent())->second;
-          assert(CV.size() == 1 && "non-unique color for block!");
-          Instruction *EHPad = CV.front()->getFirstNonPHI();
-          if (EHPad->isEHPad())
-            OpBundles.emplace_back("funclet", EHPad);
-        }
-
-        CallInst::Create(IA, None, OpBundles, "", Inst);
+        createCallInst(IA, None, "", Inst, BlockColors);
       }
     decline_rv_optimization:
       return false;
@@ -496,7 +507,7 @@ bool ObjCARCContract::tryToPeepholeInstruction(
     case ARCInstKind::Release:
       // Try to form an objc store strong from our release. If we fail, there is
       // nothing further to do below, so continue.
-      tryToContractReleaseIntoStoreStrong(Inst, Iter);
+      tryToContractReleaseIntoStoreStrong(Inst, Iter, BlockColors);
       return true;
     case ARCInstKind::User:
       // Be conservative if the function has any alloca instructions.