1 files changed, 114 insertions, 70 deletions

diff --git a/llvm/lib/CodeGen/WinEHPrepare.cpp b/llvm/lib/CodeGen/WinEHPrepare.cpp
index df2d5db8c67..8c15b7e4413 100644
--- a/llvm/lib/CodeGen/WinEHPrepare.cpp
+++ b/llvm/lib/CodeGen/WinEHPrepare.cpp
@@ -130,8 +130,9 @@ private:
                           DenseMap<BasicBlock *, Value *> &Loads, Function &F);
   void demoteNonlocalUses(Value *V, std::set<BasicBlock *> &ColorsForBB,
                           Function &F);
-  bool prepareExplicitEH(Function &F);
-  void numberFunclet(BasicBlock *InitialBB, BasicBlock *FuncletBB);
+  bool prepareExplicitEH(Function &F,
+                         SmallVectorImpl<BasicBlock *> &EntryBlocks);
+  void colorFunclets(Function &F, SmallVectorImpl<BasicBlock *> &EntryBlocks);
 
   Triple TheTriple;
 
@@ -175,6 +176,7 @@ private:
 
   std::map<BasicBlock *, std::set<BasicBlock *>> BlockColors;
   std::map<BasicBlock *, std::set<BasicBlock *>> FuncletBlocks;
+  std::map<BasicBlock *, std::set<BasicBlock *>> FuncletChildren;
 };
 
 class WinEHFrameVariableMaterializer : public ValueMaterializer {
@@ -392,20 +394,25 @@ bool WinEHPrepare::runOnFunction(Function &Fn) {
 
   SmallVector<LandingPadInst *, 4> LPads;
   SmallVector<ResumeInst *, 4> Resumes;
+  SmallVector<BasicBlock *, 4> EntryBlocks;
   bool ForExplicitEH = false;
   for (BasicBlock &BB : Fn) {
-    if (auto *LP = BB.getLandingPadInst()) {
+    Instruction *First = BB.getFirstNonPHI();
+    if (auto *LP = dyn_cast<LandingPadInst>(First)) {
       LPads.push_back(LP);
-    } else if (BB.getFirstNonPHI()->isEHPad()) {
+    } else if (First->isEHPad()) {
+      if (!ForExplicitEH)
+        EntryBlocks.push_back(&Fn.getEntryBlock());
+      if (!isa<CatchEndPadInst>(First))
+        EntryBlocks.push_back(&BB);
       ForExplicitEH = true;
-      break;
     }
     if (auto *Resume = dyn_cast<ResumeInst>(BB.getTerminator()))
       Resumes.push_back(Resume);
   }
 
   if (ForExplicitEH)
-    return prepareExplicitEH(Fn);
+    return prepareExplicitEH(Fn, EntryBlocks);
 
   // No need to prepare functions that lack landing pads.
   if (LPads.empty())
@@ -3064,72 +3071,110 @@ void llvm::calculateWinCXXEHStateNumbers(const Function *ParentFn,
     Num.processCallSite(None, ImmutableCallSite());
 }
 
-void WinEHPrepare::numberFunclet(BasicBlock *InitialBB, BasicBlock *FuncletBB) {
-  Instruction *FirstNonPHI = FuncletBB->getFirstNonPHI();
-  bool IsCatch = isa<CatchPadInst>(FirstNonPHI);
-  bool IsCleanup = isa<CleanupPadInst>(FirstNonPHI);
+void WinEHPrepare::colorFunclets(Function &F,
+                                 SmallVectorImpl<BasicBlock *> &EntryBlocks) {
+  SmallVector<std::pair<BasicBlock *, BasicBlock *>, 16> Worklist;
+  BasicBlock *EntryBlock = &F.getEntryBlock();
 
-  // Initialize the worklist with the funclet's entry point.
-  std::vector<BasicBlock *> Worklist;
-  Worklist.push_back(InitialBB);
+  // Build up the color map, which maps each block to its set of 'colors'.
+  // For any block B, the "colors" of B are the set of funclets F (possibly
+  // including a root "funclet" representing the main function), such that
+  // F will need to directly contain B or a copy of B (where the term "directly
+  // contain" is used to distinguish from being "transitively contained" in
+  // a nested funclet).
+  // Use a CFG walk driven by a worklist of (block, color) pairs.  The "color"
+  // sets attached during this processing to a block which is the entry of some
+  // funclet F is actually the set of F's parents -- i.e. the union of colors
+  // of all predecessors of F's entry.  For all other blocks, the color sets
+  // are as defined above.  A post-pass fixes up the block color map to reflect
+  // the same sense of "color" for funclet entries as for other blocks.
+
+  Worklist.push_back({EntryBlock, EntryBlock});
 
   while (!Worklist.empty()) {
-    BasicBlock *BB = Worklist.back();
-    Worklist.pop_back();
-
-    // There can be only one "pad" basic block in the funclet: the initial one.
-    if (BB != FuncletBB && BB->isEHPad())
-      continue;
-
-    // Add 'FuncletBB' as a possible color for 'BB'.
-    if (BlockColors[BB].insert(FuncletBB).second == false) {
-      // Skip basic blocks which we have already visited.
-      continue;
+    BasicBlock *Visiting;
+    BasicBlock *Color;
+    std::tie(Visiting, Color) = Worklist.pop_back_val();
+    Instruction *VisitingHead = Visiting->getFirstNonPHI();
+    if (VisitingHead->isEHPad() && !isa<CatchEndPadInst>(VisitingHead)) {
+      // Mark this as a funclet head as a member of itself.
+      FuncletBlocks[Visiting].insert(Visiting);
+      // Queue exits with the parent color.
+      for (User *Exit : VisitingHead->users()) {
+        for (BasicBlock *Succ :
+             successors(cast<Instruction>(Exit)->getParent())) {
+          if (BlockColors[Succ].insert(Color).second) {
+            Worklist.push_back({Succ, Color});
+          }
+        }
+      }
+      // Handle CatchPad specially since its successors need different colors.
+      if (CatchPadInst *CatchPad = dyn_cast<CatchPadInst>(VisitingHead)) {
+        // Visit the normal successor with the color of the new EH pad, and
+        // visit the unwind successor with the color of the parent.
+        BasicBlock *NormalSucc = CatchPad->getNormalDest();
+        if (BlockColors[NormalSucc].insert(Visiting).second) {
+          Worklist.push_back({NormalSucc, Visiting});
+        }
+        BasicBlock *UnwindSucc = CatchPad->getUnwindDest();
+        if (BlockColors[UnwindSucc].insert(Color).second) {
+          Worklist.push_back({UnwindSucc, Color});
+        }
+        continue;
+      }
+      // Switch color to the current node, except for terminate pads which
+      // have no bodies and only unwind successors and so need their successors
+      // visited with the color of the parent.
+      if (!isa<TerminatePadInst>(VisitingHead))
+        Color = Visiting;
+    } else {
+      // Note that this is a member of the given color.
+      FuncletBlocks[Color].insert(Visiting);
+      TerminatorInst *Terminator = Visiting->getTerminator();
+      if (isa<CleanupReturnInst>(Terminator) ||
+          isa<CatchReturnInst>(Terminator)) {
+        // These block's successors have already been queued with the parent
+        // color.
+        continue;
+      }
     }
+    for (BasicBlock *Succ : successors(Visiting)) {
+      if (isa<CatchEndPadInst>(Succ->getFirstNonPHI())) {
+        // The catchendpad needs to be visited with the parent's color, not
+        // the current color.  This will happen in the code above that visits
+        // any catchpad unwind successor with the parent color, so we can
+        // safely skip this successor here.
+        continue;
+      }
+      if (BlockColors[Succ].insert(Color).second) {
+        Worklist.push_back({Succ, Color});
+      }
+    }
+  }
 
-    FuncletBlocks[FuncletBB].insert(BB);
-
-    Instruction *Terminator = BB->getTerminator();
-    // The catchret's successors cannot be part of the funclet.
-    if (IsCatch && isa<CatchReturnInst>(Terminator))
-      continue;
-    // The cleanupret's successors cannot be part of the funclet.
-    if (IsCleanup && isa<CleanupReturnInst>(Terminator))
-      continue;
-
-    Worklist.insert(Worklist.end(), succ_begin(BB), succ_end(BB));
+  // The processing above actually accumulated the parent set for this
+  // funclet into the color set for its entry; use the parent set to
+  // populate the children map, and reset the color set to include just
+  // the funclet itself (no instruction can target a funclet entry except on
+  // that transitions to the child funclet).
+  for (BasicBlock *FuncletEntry : EntryBlocks) {
+    std::set<BasicBlock *> &ColorMapItem = BlockColors[FuncletEntry];
+    for (BasicBlock *Parent : ColorMapItem)
+      FuncletChildren[Parent].insert(FuncletEntry);
+    ColorMapItem.clear();
+    ColorMapItem.insert(FuncletEntry);
   }
 }
 
-bool WinEHPrepare::prepareExplicitEH(Function &F) {
+bool WinEHPrepare::prepareExplicitEH(
+    Function &F, SmallVectorImpl<BasicBlock *> &EntryBlocks) {
   // Remove unreachable blocks.  It is not valuable to assign them a color and
   // their existence can trick us into thinking values are alive when they are
   // not.
   removeUnreachableBlocks(F);
 
-  BasicBlock *EntryBlock = &F.getEntryBlock();
-
-  // Number everything starting from the entry block.
-  numberFunclet(EntryBlock, EntryBlock);
-
-  for (BasicBlock &BB : F) {
-    // Remove single entry PHIs to simplify preparation.
-    if (auto *PN = dyn_cast<PHINode>(BB.begin()))
-      if (PN->getNumIncomingValues() == 1)
-        FoldSingleEntryPHINodes(&BB);
-
-    // EH pad instructions are always the first non-PHI nodes in a block if they
-    // are at all present.
-    Instruction *I = BB.getFirstNonPHI();
-    if (I->isEHPad())
-      numberFunclet(&BB, &BB);
-
-    // It is possible for a normal basic block to only be reachable via an
-    // exceptional basic block.  The successor of a catchret is the only case
-    // where this is possible.
-    if (auto *CRI = dyn_cast<CatchReturnInst>(BB.getTerminator()))
-      numberFunclet(CRI->getSuccessor(), EntryBlock);
-  }
+  // Determine which blocks are reachable from which funclet entries.
+  colorFunclets(F, EntryBlocks);
 
   // Strip PHI nodes off of EH pads.
   SmallVector<PHINode *, 16> PHINodes;
@@ -3180,9 +3225,8 @@ bool WinEHPrepare::prepareExplicitEH(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 &Funclet : FuncletBlocks) {
-    BasicBlock *FuncletPadBB = Funclet.first;
-    std::set<BasicBlock *> &BlocksInFunclet = Funclet.second;
+  for (BasicBlock *FuncletPadBB : EntryBlocks) {
+    std::set<BasicBlock *> &BlocksInFunclet = FuncletBlocks[FuncletPadBB];
 
     std::map<BasicBlock *, BasicBlock *> Orig2Clone;
     ValueToValueMapTy VMap;
@@ -3193,12 +3237,12 @@ bool WinEHPrepare::prepareExplicitEH(Function &F) {
       if (NumColorsForBB == 1)
         continue;
 
-      assert(!isa<PHINode>(BB->front()) &&
-             "Polychromatic PHI nodes should have been demoted!");
-
       // Create a new basic block and copy instructions into it!
-      BasicBlock *CBB = CloneBasicBlock(
-          BB, VMap, Twine(".for.", FuncletPadBB->getName()), &F);
+      BasicBlock *CBB =
+          CloneBasicBlock(BB, VMap, Twine(".for.", FuncletPadBB->getName()));
+      // Insert the clone immediately after the original to ensure determinism
+      // and to keep the same relative ordering of any funclet's blocks.
+      CBB->insertInto(&F, BB->getNextNode());
 
       // Add basic block mapping.
       VMap[BB] = CBB;
@@ -3284,6 +3328,7 @@ bool WinEHPrepare::prepareExplicitEH(Function &F) {
 
   BlockColors.clear();
   FuncletBlocks.clear();
+  FuncletChildren.clear();
 
   return true;
 }
@@ -3396,12 +3441,11 @@ void WinEHPrepare::demoteNonlocalUses(Value *V,
     auto *UsingInst = cast<Instruction>(U.getUser());
     BasicBlock *UsingBB = UsingInst->getParent();
 
-    // Is the Use inside a block which is colored with a subset of the Def?
+    // Is the Use inside a block which is colored the same as the Def?
     // If so, we don't need to escape the Def because we will clone
     // ourselves our own private copy.
     std::set<BasicBlock *> &ColorsForUsingBB = BlockColors[UsingBB];
-    if (std::includes(ColorsForBB.begin(), ColorsForBB.end(),
-                      ColorsForUsingBB.begin(), ColorsForUsingBB.end()))
+    if (ColorsForUsingBB == ColorsForBB)
       continue;
 
     replaceUseWithLoad(V, U, SpillSlot, Loads, F);