1 files changed, 35 insertions, 2 deletions

diff --git a/llvm/lib/CodeGen/FuncletLayout.cpp b/llvm/lib/CodeGen/FuncletLayout.cpp
index 0cda11f53db..e1e185b4752 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 @@ static void
 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 @@ bool FuncletLayout::runOnMachineFunction(MachineFunction &F) {
   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];