[WinEH] Update catchrets with cloned successors

[JosephTremoulet]

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.

[majnemer]

I'm confused as to why this doesn't trip up our existing assertions. Preparation should have asserted because %join has two colors (none and %cleanpad). My guess is we make the verifyPreparedFunclets rerun coloring.

[majnemer]

A-ha, we generate the following IR which upsets the verifier (when running with -disable-cleanups):

join:                                             ; preds = %pad.outer, %pad.inner
  %phi = phi i32 [ 1, %pad.inner ], [ 2, %pad.outer ]
  call void @llvm.foo(i32 %phi)
  unreachable

join.for.entry:                                   ; No predecessors!
  %phi.for.entry = phi i32 
  call void @llvm.foo(i32 %phi.for.entry)
  unreachable

lib/CodeGen/WinEHPrepare.cpp

700

Can we think of a better name than BlockInFunclet? In this case, we are checking to make sure that the catchret's target makes sense given it's parent funclet.

703

I'd write this as !IncomingColors.empty().

[majnemer]

I went ahead and wrote rL256687 which results in your new test case failing regardless of how winehprepare is operated.

Your changes LGTM with nits addressed.

[JosephTremoulet]

Address feedback

lib/CodeGen/WinEHPrepare.cpp

 
 void WinEHPrepare::cloneCommonBlocks(Function &F) {
   // We need to clone all blocks which belong to multiple funclets.  Values are
   // remapped throughout the funclet to propogate both the new instructions
   // *and* the new basic blocks themselves.
   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;
     for (BasicBlock *BB : BlocksInFunclet) {
       ColorVector &ColorsForBB = BlockColors[BB];
       // We don't need to do anything if the block is monochromatic.
       size_t NumColorsForBB = ColorsForBB.size();
       if (NumColorsForBB == 1)
for (auto &Funclets : FuncletBlocks) {
     // Loop over all of the instructions in this funclet, fixing up operand
     // references as we go.  This uses VMap to do all the hard work.
     for (BasicBlock *BB : BlocksInFunclet)
       // Loop over all instructions, fixing each one as we find it...
       for (Instruction &I : *BB)
         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);
+        bool EdgeTargetsFunclet;
+        if (auto *CRI =
+                dyn_cast<CatchReturnInst>(IncomingBlock->getTerminator())) {
+          EdgeTargetsFunclet = (CRI->getParentPad() == FuncletToken);
+        } else {
           ColorVector &IncomingColors = BlockColors[IncomingBlock];
-        bool BlockInFunclet = IncomingColors.size() == 1 &&
-                              IncomingColors.front() == FuncletPadBB;
-        if (IsForOldBlock != BlockInFunclet)
+          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.
         --PredIdx;
         --PredEnd;
       }
     };
 

test/CodeGen/WinEH/wineh-cloning.ll

 ; RUN: opt -mtriple=x86_64-pc-windows-msvc -S -winehprepare  < %s | FileCheck %s
 
 declare i32 @__CxxFrameHandler3(...)
 declare i32 @__C_specific_handler(...)
+declare void @ProcessCLRException(...)
 
 declare void @f()
 
 declare void @llvm.foo(i32) nounwind
 declare void @llvm.bar() nounwind
 declare i32 @llvm.qux() nounwind
 declare i1 @llvm.baz() nounwind
 

 entry:
   ret void
 
 unreachable:
   cleanuppad within none []
   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
 ; CHECK-LABEL: !llvm.module.flags
 ; CHECK-NOT: !DISubprogram
 ; CHECK: !{{[0-9]+}} = distinct !DISubprogram(name: "test12"
 ; CHECK-NOT: !DISubprogram
 !llvm.module.flags = !{!0}