1 files changed, 25 insertions, 7 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp
index c069af9eed623..8442b49e25f45 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp
@@ -79,7 +79,6 @@ template <> bool ConcreteRegion<MachineLoop>::isLoop() const { return true; }
 class RegionInfo {
   const MachineLoopInfo &MLI;
   const WebAssemblyExceptionInfo &WEI;
-  std::vector<const Region *> Regions;
   DenseMap<const MachineLoop *, std::unique_ptr<Region>> LoopMap;
   DenseMap<const WebAssemblyException *, std::unique_ptr<Region>> ExceptionMap;
 
@@ -93,7 +92,14 @@ public:
     const auto *WE = WEI.getExceptionFor(MBB);
     if (!ML && !WE)
       return nullptr;
-    if ((ML && !WE) || (ML && WE && ML->getNumBlocks() < WE->getNumBlocks())) {
+    // We determine subregion relationship by domination of their headers, i.e.,
+    // if region A's header dominates region B's header, B is a subregion of A.
+    // WebAssemblyException contains BBs in all its subregions (loops or
+    // exceptions), but MachineLoop may not, because MachineLoop does not contain
+    // BBs that don't have a path to its header even if they are dominated by
+    // its header. So here we should use WE->contains(ML->getHeader()), but not
+    // ML->contains(WE->getHeader()).
+    if ((ML && !WE) || (ML && WE && WE->contains(ML->getHeader()))) {
       // If the smallest region containing MBB is a loop
       if (LoopMap.count(ML))
         return LoopMap[ML].get();
@@ -152,9 +158,17 @@ static void maybeUpdateTerminator(MachineBasicBlock *MBB) {
     AllAnalyzable &= Term.isBranch() && !Term.isIndirectBranch();
   }
   assert((AnyBarrier || AllAnalyzable) &&
-         "AnalyzeBranch needs to analyze any block with a fallthrough");
+         "analyzeBranch needs to analyze any block with a fallthrough");
+
+  // Find the layout successor from the original block order.
+  MachineFunction *MF = MBB->getParent();
+  MachineBasicBlock *OriginalSuccessor =
+      unsigned(MBB->getNumber() + 1) < MF->getNumBlockIDs()
+          ? MF->getBlockNumbered(MBB->getNumber() + 1)
+          : nullptr;
+
   if (AllAnalyzable)
-    MBB->updateTerminator();
+    MBB->updateTerminator(OriginalSuccessor);
 }
 
 namespace {
@@ -241,9 +255,12 @@ struct Entry {
 static void sortBlocks(MachineFunction &MF, const MachineLoopInfo &MLI,
                        const WebAssemblyExceptionInfo &WEI,
                        const MachineDominatorTree &MDT) {
+  // Remember original layout ordering, so we can update terminators after
+  // reordering to point to the original layout successor.
+  MF.RenumberBlocks();
+
   // Prepare for a topological sort: Record the number of predecessors each
   // block has, ignoring loop backedges.
-  MF.RenumberBlocks();
   SmallVector<unsigned, 16> NumPredsLeft(MF.getNumBlockIDs(), 0);
   for (MachineBasicBlock &MBB : MF) {
     unsigned N = MBB.pred_size();
@@ -368,6 +385,7 @@ static void sortBlocks(MachineFunction &MF, const MachineLoopInfo &MLI,
     const Region *Region = RI.getRegionFor(&MBB);
 
     if (Region && &MBB == Region->getHeader()) {
+      // Region header.
       if (Region->isLoop()) {
         // Loop header. The loop predecessor should be sorted above, and the
         // other predecessors should be backedges below.
@@ -377,7 +395,7 @@ static void sortBlocks(MachineFunction &MF, const MachineLoopInfo &MLI,
               "Loop header predecessors must be loop predecessors or "
               "backedges");
       } else {
-        // Not a loop header. All predecessors should be sorted above.
+        // Exception header. All predecessors should be sorted above.
         for (auto Pred : MBB.predecessors())
           assert(Pred->getNumber() < MBB.getNumber() &&
                  "Non-loop-header predecessors should be topologically sorted");
@@ -386,7 +404,7 @@ static void sortBlocks(MachineFunction &MF, const MachineLoopInfo &MLI,
              "Regions should be declared at most once.");
 
     } else {
-      // Not a loop header. All predecessors should be sorted above.
+      // Not a region header. All predecessors should be sorted above.
       for (auto Pred : MBB.predecessors())
         assert(Pred->getNumber() < MBB.getNumber() &&
                "Non-loop-header predecessors should be topologically sorted");