[ObjCARC] Account for funclet token in storeStrong transform

[smeenai]

Summary

When creating a call to storeStrong in ObjCARCContract, ensure the call
gets the correct funclet token, otherwise WinEHPrepare will turn the
call (and all subsequent instructions) into unreachable.

We already have logic to do this for the ARC autorelease elision marker;
factor that out into a common function that's used for both. These are
the only two places in this transform that create call instructions.

[rnk]

lgtm

I'm unhappy with the need for instrumentation passes to "color" their inserted call instructions. It's a real drag. What we really want here are single-entry, multi-exit regions. If you ignore unreachable, then funclets are already single-entry, single-exit regions. Unfortunately, unreachable is very very common, because it comes after throw, which always calls a noreturn function. So this is common:

try { may_throw(); } catch (...) { do_something(); throw; /*unreachable*/ }
try { may_throw(); } catch (...) { do_something(); throw; /*unreachable*/ }

We needed to prevent simplifycfg from tail merging the two otherwise identical catch funclets that don't end in catchret.

If we could require instead that those unreachable instructions refer to their parent funclet, we'd be made in the shade. WinEHPrepare wouldn't have to do BB cloning anymore.

Throw it on the TODO list, I guess.

lib/Transforms/ObjCARC/ObjCARCContract.cpp

class ObjCARCContract : public FunctionPass {
     bool optimizeRetainCall(Function &F, Instruction *Retain);
 
     bool
     contractAutorelease(Function &F, Instruction *Autorelease,
                         ARCInstKind Class,
                         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;
     bool runOnFunction(Function &F) override;
 
   public:
     static char ID;
     ObjCARCContract() : FunctionPass(ID) {
findRetainForStoreStrongContraction(Value *New, StoreInst *Store,
   Instruction *Retain = &*I;
   if (GetBasicARCInstKind(Retain) != ARCInstKind::Retain)
     return nullptr;
   if (GetArgRCIdentityRoot(Retain) != New)
     return nullptr;
   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:
 ///
 ///   objc_storeStrong(i8** %old_ptr, i8* new_value)
 ///
 /// is equivalent to the following IR sequence:
 ///
 ///   ; Load old value.

 /// The safety of this optimization is based around the following
 /// considerations:
 ///
 ///  1. We are forming the store strong at the store. Thus to perform this
 ///     optimization it must be safe to move the retain, load, and release to
 ///     (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())
     return;
 
   // For now, require everything to be in one basic block.
   BasicBlock *BB = Release->getParent();
   if (Load->getParent() != BB)
void ObjCARCContract::tryToContractReleaseIntoStoreStrong(
   Type *I8XX = PointerType::getUnqual(I8X);
 
   Value *Args[] = { Load->getPointerOperand(), New };
   if (Args[0]->getType() != I8XX)
     Args[0] = new BitCastInst(Args[0], I8XX, "", Store);
   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());
 
   // We can't set the tail flag yet, because we haven't yet determined
   // whether there are any escaping allocas. Remember this call, so that
   // we can set the tail flag once we know it's safe.
   StoreStrongCalls.insert(StoreStrong);
 
case ARCInstKind::ClaimRV: {
                         "optimization.\n");
         Changed = true;
         InlineAsm *IA = InlineAsm::get(
             FunctionType::get(Type::getVoidTy(Inst->getContext()),
                               /*isVarArg=*/false),
             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;
     }
     case ARCInstKind::InitWeak: {
       // objc_initWeak(p, null) => *p = null
       CallInst *CI = cast<CallInst>(Inst);
       if (IsNullOrUndef(CI->getArgOperand(1))) {
         Value *Null =
           ConstantPointerNull::get(cast<PointerType>(CI->getType()));
         Changed = true;
         new StoreInst(Null, CI->getArgOperand(0), CI);
 
         DEBUG(dbgs() << "OBJCARCContract: Old = " << *CI << "\n"
                      << "                 New = " << *Null << "\n");
 
         CI->replaceAllUsesWith(Null);
         CI->eraseFromParent();
       }
       return true;
     }
     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.
       // Technically we only care about escaping alloca instructions,
       // but this is sufficient to handle some interesting cases.
       if (isa<AllocaInst>(Inst))
         TailOkForStoreStrongs = false;
       return true;

test/Transforms/ObjCARC/contract-storestrong-funclet.ll

+; RUN: opt -mtriple=i686-unknown-windows-msvc -objc-arc-contract -S -o - %s | FileCheck %s
+
+declare void @f()
+declare i32 @__CxxFrameHandler3(...)
+declare dllimport i8* @objc_retain(i8*)
+declare dllimport i8* @objc_retainAutoreleasedReturnValue(i8*)
+declare dllimport void @objc_release(i8*)
+
+@x = external global i8*
+
+define void @g(i8* %p) personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
+  invoke void @f() to label %invoke.cont unwind label %ehcleanup
+
+invoke.cont:
+  %call = tail call i8* @objc_retain(i8* %p) nounwind
+  %tmp = load i8*, i8** @x, align 4
+  store i8* %call, i8** @x, align 4
+  tail call void @objc_release(i8* %tmp) nounwind
+  ret void
+
+ehcleanup:
+  %1 = cleanuppad within none []
+  %call1 = tail call i8* @objc_retain(i8* %p) nounwind [ "funclet"(token %1) ]
+  %tmp1 = load i8*, i8** @x, align 4
+  store i8* %call1, i8** @x, align 4
+  tail call void @objc_release(i8* %tmp1) nounwind [ "funclet"(token %1) ]
+  cleanupret from %1 unwind to caller
+}
+
+; CHECK-LABEL: invoke.cont:
+; CHECK: tail call void @objc_storeStrong(i8** @x, i8* %p) #0{{$}}
+; CHECK: ret void
+
+; CHECK-LABEL: ehcleanup:
+; CHECK: %1 = cleanuppad within none []
+; CHECK: tail call void @objc_storeStrong(i8** @x, i8* %p) #0 [ "funclet"(token %1) ]
+; CHECK: cleanupret from %1 unwind to caller