[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];
diff --git a/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp b/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
index baf9ee9..7da68de 100644
--- a/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
@@ -290,9 +290,11 @@
   const Function *WinEHParentFn = MMI.getWinEHParent(&fn);
   if (Personality == EHPersonality::MSVC_CXX)
     calculateWinCXXEHStateNumbers(WinEHParentFn, EHInfo);
-  else
+  else if (isAsynchronousEHPersonality(Personality))
     calculateSEHStateNumbers(WinEHParentFn, EHInfo);
 
+  calculateCatchReturnSuccessorColors(WinEHParentFn, EHInfo);
+
   // Map all BB references in the WinEH data to MBBs.
   for (WinEHTryBlockMapEntry &TBME : EHInfo.TryBlockMap) {
     for (WinEHHandlerType &H : TBME.HandlerArray) {
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index 4f96565..f779fd1b 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -1168,9 +1168,21 @@
   MachineBasicBlock *TargetMBB = FuncInfo.MBBMap[I.getSuccessor()];
   FuncInfo.MBB->addSuccessor(TargetMBB);
 
+  // Figure out the funclet membership for the catchret's successor.
+  // This will be used by the FuncletLayout pass to determine how to order the
+  // BB's.
+  MachineModuleInfo &MMI = DAG.getMachineFunction().getMMI();
+  WinEHFuncInfo &EHInfo =
+      MMI.getWinEHFuncInfo(DAG.getMachineFunction().getFunction());
+  const BasicBlock *SuccessorColor = EHInfo.CatchRetSuccessorColorMap[&I];
+  assert(SuccessorColor && "No parent funclet for catchret!");
+  MachineBasicBlock *SuccessorColorMBB = FuncInfo.MBBMap[SuccessorColor];
+  assert(SuccessorColorMBB && "No MBB for SuccessorColor!");
+
   // Create the terminator node.
   SDValue Ret = DAG.getNode(ISD::CATCHRET, getCurSDLoc(), MVT::Other,
-                            getControlRoot(), DAG.getBasicBlock(TargetMBB));
+                            getControlRoot(), DAG.getBasicBlock(TargetMBB),
+                            DAG.getBasicBlock(SuccessorColorMBB));
   DAG.setRoot(Ret);
 }
 
diff --git a/llvm/lib/CodeGen/WinEHPrepare.cpp b/llvm/lib/CodeGen/WinEHPrepare.cpp
index 7240159..07734c3 100644
--- a/llvm/lib/CodeGen/WinEHPrepare.cpp
+++ b/llvm/lib/CodeGen/WinEHPrepare.cpp
@@ -399,6 +399,29 @@
   return new WinEHPrepare(TM);
 }
 
+static bool
+findExceptionalConstructs(Function &Fn,
+                          SmallVectorImpl<LandingPadInst *> &LPads,
+                          SmallVectorImpl<ResumeInst *> &Resumes,
+                          SmallVectorImpl<BasicBlock *> &EntryBlocks) {
+  bool ForExplicitEH = false;
+  for (BasicBlock &BB : Fn) {
+    Instruction *First = BB.getFirstNonPHI();
+    if (auto *LP = dyn_cast<LandingPadInst>(First)) {
+      LPads.push_back(LP);
+    } else if (First->isEHPad()) {
+      if (!ForExplicitEH)
+        EntryBlocks.push_back(&Fn.getEntryBlock());
+      if (!isa<CatchEndPadInst>(First) && !isa<CleanupEndPadInst>(First))
+        EntryBlocks.push_back(&BB);
+      ForExplicitEH = true;
+    }
+    if (auto *Resume = dyn_cast<ResumeInst>(BB.getTerminator()))
+      Resumes.push_back(Resume);
+  }
+  return ForExplicitEH;
+}
+
 bool WinEHPrepare::runOnFunction(Function &Fn) {
   if (!Fn.hasPersonalityFn())
     return false;
@@ -417,21 +440,8 @@
   SmallVector<LandingPadInst *, 4> LPads;
   SmallVector<ResumeInst *, 4> Resumes;
   SmallVector<BasicBlock *, 4> EntryBlocks;
-  bool ForExplicitEH = false;
-  for (BasicBlock &BB : Fn) {
-    Instruction *First = BB.getFirstNonPHI();
-    if (auto *LP = dyn_cast<LandingPadInst>(First)) {
-      LPads.push_back(LP);
-    } else if (First->isEHPad()) {
-      if (!ForExplicitEH)
-        EntryBlocks.push_back(&Fn.getEntryBlock());
-      if (!isa<CatchEndPadInst>(First) && !isa<CleanupEndPadInst>(First))
-        EntryBlocks.push_back(&BB);
-      ForExplicitEH = true;
-    }
-    if (auto *Resume = dyn_cast<ResumeInst>(BB.getTerminator()))
-      Resumes.push_back(Resume);
-  }
+  bool ForExplicitEH =
+      findExceptionalConstructs(Fn, LPads, Resumes, EntryBlocks);
 
   if (ForExplicitEH)
     return prepareExplicitEH(Fn, EntryBlocks);
@@ -2878,8 +2888,11 @@
   }
 }
 
-void WinEHPrepare::colorFunclets(Function &F,
-                                 SmallVectorImpl<BasicBlock *> &EntryBlocks) {
+static void
+colorFunclets(Function &F, SmallVectorImpl<BasicBlock *> &EntryBlocks,
+              std::map<BasicBlock *, std::set<BasicBlock *>> &BlockColors,
+              std::map<BasicBlock *, std::set<BasicBlock *>> &FuncletBlocks,
+              std::map<BasicBlock *, std::set<BasicBlock *>> &FuncletChildren) {
   SmallVector<std::pair<BasicBlock *, BasicBlock *>, 16> Worklist;
   BasicBlock *EntryBlock = &F.getEntryBlock();
 
@@ -2976,6 +2989,56 @@
   }
 }
 
+void WinEHPrepare::colorFunclets(Function &F,
+                                 SmallVectorImpl<BasicBlock *> &EntryBlocks) {
+  ::colorFunclets(F, EntryBlocks, BlockColors, FuncletBlocks, FuncletChildren);
+}
+
+void llvm::calculateCatchReturnSuccessorColors(const Function *Fn,
+                                               WinEHFuncInfo &FuncInfo) {
+  SmallVector<LandingPadInst *, 4> LPads;
+  SmallVector<ResumeInst *, 4> Resumes;
+  SmallVector<BasicBlock *, 4> EntryBlocks;
+  // colorFunclets needs the set of EntryBlocks, get them using
+  // findExceptionalConstructs.
+  bool ForExplicitEH = findExceptionalConstructs(const_cast<Function &>(*Fn),
+                                                 LPads, Resumes, EntryBlocks);
+  if (!ForExplicitEH)
+    return;
+
+  std::map<BasicBlock *, std::set<BasicBlock *>> BlockColors;
+  std::map<BasicBlock *, std::set<BasicBlock *>> FuncletBlocks;
+  std::map<BasicBlock *, std::set<BasicBlock *>> FuncletChildren;
+  // Figure out which basic blocks belong to which funclets.
+  colorFunclets(const_cast<Function &>(*Fn), EntryBlocks, BlockColors,
+                FuncletBlocks, FuncletChildren);
+
+  // We need to find the catchret successors.  To do this, we must first find
+  // all the catchpad funclets.
+  for (auto &Funclet : FuncletBlocks) {
+    // Figure out what kind of funclet we are looking at; We only care about
+    // catchpads.
+    BasicBlock *FuncletPadBB = Funclet.first;
+    Instruction *FirstNonPHI = FuncletPadBB->getFirstNonPHI();
+    auto *CatchPad = dyn_cast<CatchPadInst>(FirstNonPHI);
+    if (!CatchPad)
+      continue;
+
+    // The users of a catchpad are always catchrets.
+    for (User *Exit : CatchPad->users()) {
+      auto *CatchReturn = cast<CatchReturnInst>(Exit);
+      BasicBlock *CatchRetSuccessor = CatchReturn->getSuccessor();
+      std::set<BasicBlock *> &SuccessorColors = BlockColors[CatchRetSuccessor];
+      assert(SuccessorColors.size() == 1 && "Expected BB to be monochrome!");
+      BasicBlock *Color = *SuccessorColors.begin();
+      if (auto *CPI = dyn_cast<CatchPadInst>(Color->getFirstNonPHI()))
+        Color = CPI->getNormalDest();
+      // Record the catchret successor's funclet membership.
+      FuncInfo.CatchRetSuccessorColorMap[CatchReturn] = Color;
+    }
+  }
+}
+
 void WinEHPrepare::demotePHIsOnFunclets(Function &F) {
   // Strip PHI nodes off of EH pads.
   SmallVector<PHINode *, 16> PHINodes;