[WinEH] Update catchrets with cloned successors

Summary:
Add a pass to update catchrets when their successors get cloned; the
existing pass doesn't catch these because it walks the funclet whose
blocks are being cloned but the catchret is in a child funclet.

Also update the test for removing incoming PHI values; when the
predecessor is a catchret, the relevant color is the catchret's parentPad,
not its block's color.


Reviewers: andrew.w.kaylor, rnk, majnemer

Subscribers: llvm-commits

Differential Revision: http://reviews.llvm.org/D15840

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@256689 91177308-0d34-0410-b5e6-96231b3b80d8

lib/CodeGen/WinEHPrepare.cpp

@@ -598,6 +598,11 @@ void WinEHPrepare::cloneCommonBlocks(Function &F) {
   for (auto &Funclets : FuncletBlocks) {
     BasicBlock *FuncletPadBB = Funclets.first;
     std::vector<BasicBlock *> &BlocksInFunclet = Funclets.second;
+    Value *FuncletToken;
+    if (FuncletPadBB == &F.getEntryBlock())
+      FuncletToken = ConstantTokenNone::get(F.getContext());
+    else
+      FuncletToken = FuncletPadBB->getFirstNonPHI();
 
     std::vector<std::pair<BasicBlock *, BasicBlock *>> Orig2Clone;
     ValueToValueMapTy VMap;
@@ -669,15 +674,44 @@ void WinEHPrepare::cloneCommonBlocks(Function &F) {
         RemapInstruction(&I, VMap,
                          RF_IgnoreMissingEntries | RF_NoModuleLevelChanges);
 
+    // Catchrets targeting cloned blocks need to be updated separately from
+    // the loop above because they are not in the current funclet.
+    SmallVector<CatchReturnInst *, 2> FixupCatchrets;
+    for (auto &BBMapping : Orig2Clone) {
+      BasicBlock *OldBlock = BBMapping.first;
+      BasicBlock *NewBlock = BBMapping.second;
+
+      FixupCatchrets.clear();
+      for (BasicBlock *Pred : predecessors(OldBlock))
+        if (auto *CatchRet = dyn_cast<CatchReturnInst>(Pred->getTerminator()))
+          if (CatchRet->getParentPad() == FuncletToken)
+            FixupCatchrets.push_back(CatchRet);
+
+      for (CatchReturnInst *CatchRet : FixupCatchrets)
+        CatchRet->setSuccessor(NewBlock);
+    }
+
     auto UpdatePHIOnClonedBlock = [&](PHINode *PN, bool IsForOldBlock) {
       unsigned NumPreds = PN->getNumIncomingValues();
       for (unsigned PredIdx = 0, PredEnd = NumPreds; PredIdx != PredEnd;
            ++PredIdx) {
         BasicBlock *IncomingBlock = PN->getIncomingBlock(PredIdx);
-        ColorVector &IncomingColors = BlockColors[IncomingBlock];
-        bool BlockInFunclet = IncomingColors.size() == 1 &&
-                              IncomingColors.front() == FuncletPadBB;
-        if (IsForOldBlock != BlockInFunclet)
+        bool EdgeTargetsFunclet;
+        if (auto *CRI =
+                dyn_cast<CatchReturnInst>(IncomingBlock->getTerminator())) {
+          EdgeTargetsFunclet = (CRI->getParentPad() == FuncletToken);
+        } else {
+          ColorVector &IncomingColors = BlockColors[IncomingBlock];
+          assert(!IncomingColors.empty() && "Block not colored!");
+          assert((IncomingColors.size() == 1 ||
+                  llvm::all_of(IncomingColors,
+                               [&](BasicBlock *Color) {
+                                 return Color != FuncletPadBB;
+                               })) &&
+                 "Cloning should leave this funclet's blocks monochromatic");
+          EdgeTargetsFunclet = (IncomingColors.front() == FuncletPadBB);
+        }
+        if (IsForOldBlock != EdgeTargetsFunclet)
           continue;
         PN->removeIncomingValue(IncomingBlock, /*DeletePHIIfEmpty=*/false);
         // Revisit the next entry.

test/CodeGen/WinEH/wineh-cloning.ll

@@ -2,6 +2,7 @@
 
 declare i32 @__CxxFrameHandler3(...)
 declare i32 @__C_specific_handler(...)
+declare void @ProcessCLRException(...)
 
 declare void @f()
 
@@ -369,6 +370,50 @@ unreachable:
   unreachable
 }
 
+define void @test14() personality void (...)* @ProcessCLRException {
+entry:
+  invoke void @f()
+    to label %cont unwind label %cleanup
+cont:
+  invoke void @f()
+    to label %exit unwind label %switch.outer
+cleanup:
+  %cleanpad = cleanuppad within none []
+  invoke void @f() [ "funclet" (token %cleanpad) ]
+    to label %cleanret unwind label %switch.inner
+switch.inner:
+  %cs.inner = catchswitch within %cleanpad [label %pad.inner] unwind to caller
+pad.inner:
+  %cp.inner = catchpad within %cs.inner [i32 1]
+  catchret from %cp.inner to label %join
+cleanret:
+  cleanupret from %cleanpad unwind to caller
+switch.outer:
+  %cs.outer = catchswitch within none [label %pad.outer] unwind to caller
+pad.outer:
+  %cp.outer = catchpad within %cs.outer [i32 2]
+  catchret from %cp.outer to label %join
+join:
+  %phi = phi i32 [ 1, %pad.inner ], [ 2, %pad.outer ]
+  call void @llvm.foo(i32 %phi)
+  unreachable
+exit:
+  ret void
+}
+; Both catchrets target %join, but the catchret from %cp.inner
+; returns to %cleanpad and the catchret from %cp.outer returns to the
+; main function, so %join needs to get cloned and one of the cleanuprets
+; needs to be updated to target the clone
+; CHECK-LABEL: define void @test14()
+; CHECK: catchret from %cp.inner to label %[[Clone1:.+]]
+; CHECK: catchret from %cp.outer to label %[[Clone2:.+]]
+; CHECK: [[Clone1]]:
+; CHECK-NEXT: call void @llvm.foo(i32 1)
+; CHECK-NEXT: unreachable
+; CHECK: [[Clone2]]:
+; CHECK-NEXT: call void @llvm.foo(i32 2)
+; CHECK-NEXT: unreachable
+
 ;; Debug info (from test12)
 
 ; Make sure the DISubprogram doesn't get cloned