[WinEH] Update coloring to handle nested cases cleanly

Summary:
Change the coloring algorithm in WinEHPrepare to visit a funclet's exits
in its parents' contexts and so properly classify the continuations of
nested funclets.

Also change the placement of cloned blocks to be deterministic and to
maintain the relative order of each funclet's blocks.

Add a lit test showing various patterns that require cloning, the last
several of which don't have CHECKs yet because they require cloning
entire funclets which is NYI.


Reviewers: rnk, andrew.w.kaylor, majnemer

Subscribers: llvm-commits

Differential Revision: http://reviews.llvm.org/D12353

llvm-svn: 246245
diff --git a/llvm/lib/CodeGen/WinEHPrepare.cpp b/llvm/lib/CodeGen/WinEHPrepare.cpp
index df2d5db..8c15b7e 100644
--- a/llvm/lib/CodeGen/WinEHPrepare.cpp
+++ b/llvm/lib/CodeGen/WinEHPrepare.cpp
@@ -130,8 +130,9 @@
                           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 @@
 
   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 @@
 
   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 @@
     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 @@
   // 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 @@
       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 @@
 
   BlockColors.clear();
   FuncletBlocks.clear();
+  FuncletChildren.clear();
 
   return true;
 }
@@ -3396,12 +3441,11 @@
     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);