[WinEH] Make FuncletLayout more robust against catchret

Catchret transfers control from a catch funclet to an earlier funclet.
However, it is not completely clear which funclet the catchret target is
part of.  Make this clear by stapling the catchret target's funclet
membership onto the CATCHRET SDAG node.

llvm-svn: 249052
diff --git a/llvm/lib/CodeGen/FuncletLayout.cpp b/llvm/lib/CodeGen/FuncletLayout.cpp
index 0cda11f..e1e185b 100644
--- a/llvm/lib/CodeGen/FuncletLayout.cpp
+++ b/llvm/lib/CodeGen/FuncletLayout.cpp
@@ -16,6 +16,8 @@
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/Target/TargetInstrInfo.h"
+#include "llvm/Target/TargetSubtargetInfo.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "funclet-layout"
@@ -36,8 +38,21 @@
 collectFuncletMembers(DenseMap<MachineBasicBlock *, int> &FuncletMembership,
                       int Funclet, MachineBasicBlock *MBB) {
   // Don't revisit blocks.
-  if (FuncletMembership.count(MBB) > 0)
+  if (FuncletMembership.count(MBB) > 0) {
+    // FIXME: This is a hack, we need to assert this unconditionally.
+    bool IsProbablyUnreachableBlock =
+        MBB->empty() ||
+        (MBB->succ_empty() && !MBB->getFirstTerminator()->isReturn() &&
+         MBB->size() == 1);
+
+    if (!IsProbablyUnreachableBlock) {
+      if (FuncletMembership[MBB] != Funclet) {
+        assert(false && "MBB is part of two funclets!");
+        report_fatal_error("MBB is part of two funclets!");
+      }
+    }
     return;
+  }
 
   // Add this MBB to our funclet.
   FuncletMembership[MBB] = Funclet;
@@ -71,18 +86,36 @@
   if (!F.getMMI().hasEHFunclets())
     return false;
 
+  const TargetInstrInfo *TII = F.getSubtarget().getInstrInfo();
   SmallVector<MachineBasicBlock *, 16> FuncletBlocks;
-  for (MachineBasicBlock &MBB : F)
+  SmallVector<std::pair<MachineBasicBlock *, int>, 16> CatchRetSuccessors;
+  for (MachineBasicBlock &MBB : F) {
     if (MBB.isEHFuncletEntry())
       FuncletBlocks.push_back(&MBB);
 
+    MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator();
+    if (MBBI->getOpcode() != TII->getCatchReturnOpcode())
+      continue;
+
+    MachineBasicBlock *Successor = MBBI->getOperand(0).getMBB();
+    MachineBasicBlock *SuccessorColor = MBBI->getOperand(1).getMBB();
+    CatchRetSuccessors.push_back({Successor, SuccessorColor->getNumber()});
+  }
+
   // We don't have anything to do if there aren't any EH pads.
   if (FuncletBlocks.empty())
     return false;
 
   DenseMap<MachineBasicBlock *, int> FuncletMembership;
+  // Identify all the basic blocks reachable from the function entry.
+  collectFuncletMembers(FuncletMembership, F.front().getNumber(), F.begin());
+  // Next, identify all the blocks inside the funclets.
   for (MachineBasicBlock *MBB : FuncletBlocks)
     collectFuncletMembers(FuncletMembership, MBB->getNumber(), MBB);
+  // Finally, identify all the targets of a catchret.
+  for (std::pair<MachineBasicBlock *, int> CatchRetPair : CatchRetSuccessors)
+    collectFuncletMembers(FuncletMembership, CatchRetPair.second,
+                          CatchRetPair.first);
 
   F.sort([&](MachineBasicBlock &x, MachineBasicBlock &y) {
     return FuncletMembership[&x] < FuncletMembership[&y];