[WebAssembly] Add unreachable before catch destinations (#123915)

When `try_table`'s catch clause's destination has a return type, as in
the case of catch with a concrete tag, catch_ref, and catch_all_ref. For
example:
```wasm
block exnref
  try_table (catch_all_ref 0)
    ...
  end_try_table
end_block
... use exnref ...
```

This code is not valid because the block's body type is not exnref. So
we add an unreachable after the 'end_try_table' to make the code valid
here:
```wasm
block exnref
  try_table (catch_all_ref 0)
    ...
  end_try_table
  unreachable                    ;; Newly added
end_block
```
Because 'unreachable' is a terminator we also need to split the BB.

---

We need to handle the same thing for unwind mismatch handling. In the
code below, we create a "trampoline BB" that will be the destination for
the nested `try_table`~`end_try_table` added to fix a unwind mismatch:
```wasm
try_table (catch ... )
  block exnref
    ...
    try_table (catch_all_ref N)
      some code
    end_try_table
    ...
  end_block                      ;; Trampoline BB
  throw_ref
end_try_table
```
While the `block` added for the trampoline BB has the return type
`exnref`, its body, which contains the nested `try_table` and other
code, wouldn't have the `exnref` return type. Most times it didn't
become a problem because the block's body ended with something like `br`
or `return`, but that may not always be the case, especially when there
is a loop. So we add an `unreachable` to make the code valid here too:
```wasm
try_table (catch ... )
  block exnref
    ...
    try_table (catch_all_ref N)
      some code
    end_try_table
    ...
    unreachable                  ;; Newly added
  end_block                      ;; Trampoline BB
  throw_ref
end_try_table
```
In this case we just append the `unreachable` at the end of the layout
predecessor BB. (This was tricky to do in the first (non-mismatch) case
because there `end_try_table` and `end_block` were added in the
beginning of an EH pad in `placeTryTableMarker` and moving
`end_try_table` and the new `unreachable` to the previous BB caused
other problems.)

---

This adds many `unreaachable`s to the output, but this adds
`unreachable` to only a few places to see if this is working. The
FileCheck lines in `exception.ll` and `cfg-stackify-eh.ll` are already
heavily redacted to only leave important control-flow instructions, so I
don't think it's worth adding `unreachable`s everywhere.

1 files changed, 88 insertions, 19 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
index 6cae0e7..bdc1cc6 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
@@ -1297,6 +1297,7 @@ void WebAssemblyCFGStackify::addNestedTryDelegate(
 //       some code
 //     end_try_table
 //     ...
+//     unreachable
 //   end_block                      ;; Trampoline BB
 //   throw_ref
 // end_try_table
@@ -1358,6 +1359,13 @@ WebAssemblyCFGStackify::getTrampolineBlock(MachineBasicBlock *UnwindDest) {
   BuildMI(TrampolineBB, EndDebugLoc, TII.get(WebAssembly::THROW_REF))
       .addReg(ExnReg);
 
+  // The trampoline BB's return type is exnref because it is a target of
+  // catch_all_ref. But the body type of the block we just created is not. We
+  // add an 'unreachable' right before the 'end_block' to make the code valid.
+  MachineBasicBlock *TrampolineLayoutPred = TrampolineBB->getPrevNode();
+  BuildMI(TrampolineLayoutPred, TrampolineLayoutPred->findBranchDebugLoc(),
+          TII.get(WebAssembly::UNREACHABLE));
+
   registerScope(Block, EndBlock);
   UnwindDestToTrampoline[UnwindDest] = TrampolineBB;
   return TrampolineBB;
@@ -1465,7 +1473,7 @@ void WebAssemblyCFGStackify::addNestedTryTable(MachineInstr *RangeBegin,
       // - After:
       // pre_bb: (new)
       //   range_end
-      // end_try_table: (new)
+      // end_try_table_bb: (new)
       //   end_try_table
       // post_bb: (previous 'ehpad')
       //   catch
@@ -1523,9 +1531,9 @@ void WebAssemblyCFGStackify::addNestedTryTable(MachineInstr *RangeBegin,
 //   end_loop
 //   end_try_table
 //
-// So if the unwind dest BB has a end_loop before an end_try_table, we split the
-// BB with the end_loop as a separate BB before the end_try_table BB, so that
-// after we fix the unwind mismatch, the code will be like:
+// So if an end_try_table BB has an end_loop before the end_try_table, we split
+// the BB with the end_loop as a separate BB before the end_try_table BB, so
+// that after we fix the unwind mismatch, the code will be like:
 // bb0:
 //   try_table
 //   block exnref
@@ -1538,10 +1546,10 @@ void WebAssemblyCFGStackify::addNestedTryTable(MachineInstr *RangeBegin,
 //   end_block
 // end_try_table_bb:
 //   end_try_table
-static void splitEndLoopBB(MachineBasicBlock *UnwindDest) {
-  auto &MF = *UnwindDest->getParent();
+static void splitEndLoopBB(MachineBasicBlock *EndTryTableBB) {
+  auto &MF = *EndTryTableBB->getParent();
   MachineInstr *EndTryTable = nullptr, *EndLoop = nullptr;
-  for (auto &MI : reverse(*UnwindDest)) {
+  for (auto &MI : reverse(*EndTryTableBB)) {
     if (MI.getOpcode() == WebAssembly::END_TRY_TABLE) {
       EndTryTable = &MI;
       continue;
@@ -1555,11 +1563,11 @@ static void splitEndLoopBB(MachineBasicBlock *UnwindDest) {
     return;
 
   auto *EndLoopBB = MF.CreateMachineBasicBlock();
-  MF.insert(UnwindDest->getIterator(), EndLoopBB);
+  MF.insert(EndTryTableBB->getIterator(), EndLoopBB);
   auto SplitPos = std::next(EndLoop->getIterator());
-  EndLoopBB->splice(EndLoopBB->end(), UnwindDest, UnwindDest->begin(),
+  EndLoopBB->splice(EndLoopBB->end(), EndTryTableBB, EndTryTableBB->begin(),
                     SplitPos);
-  EndLoopBB->addSuccessor(UnwindDest);
+  EndLoopBB->addSuccessor(EndTryTableBB);
 }
 
 bool WebAssemblyCFGStackify::fixCallUnwindMismatches(MachineFunction &MF) {
@@ -1943,8 +1951,16 @@ bool WebAssemblyCFGStackify::fixCallUnwindMismatches(MachineFunction &MF) {
   // When end_loop is before end_try_table within the same BB in unwind
   // destinations, we should split the end_loop into another BB.
   if (!WebAssembly::WasmUseLegacyEH)
-    for (auto &[UnwindDest, _] : UnwindDestToTryRanges)
-      splitEndLoopBB(UnwindDest);
+    for (auto &[UnwindDest, _] : UnwindDestToTryRanges) {
+      auto It = EHPadToTry.find(UnwindDest);
+      // If UnwindDest is the fake caller block, it will not be in EHPadToTry
+      // map
+      if (It != EHPadToTry.end()) {
+        auto *TryTable = It->second;
+        auto *EndTryTable = BeginToEnd[TryTable];
+        splitEndLoopBB(EndTryTable->getParent());
+      }
+    }
 
   // Now we fix the mismatches by wrapping calls with inner try-delegates.
   for (auto &P : UnwindDestToTryRanges) {
@@ -2179,8 +2195,15 @@ bool WebAssemblyCFGStackify::fixCatchUnwindMismatches(MachineFunction &MF) {
 
   // When end_loop is before end_try_table within the same BB in unwind
   // destinations, we should split the end_loop into another BB.
-  for (auto &[_, UnwindDest] : EHPadToUnwindDest)
-    splitEndLoopBB(UnwindDest);
+  for (auto &[_, UnwindDest] : EHPadToUnwindDest) {
+    auto It = EHPadToTry.find(UnwindDest);
+    // If UnwindDest is the fake caller block, it will not be in EHPadToTry map
+    if (It != EHPadToTry.end()) {
+      auto *TryTable = It->second;
+      auto *EndTryTable = BeginToEnd[TryTable];
+      splitEndLoopBB(EndTryTable->getParent());
+    }
+  }
 
   NumCatchUnwindMismatches += EHPadToUnwindDest.size();
   SmallPtrSet<MachineBasicBlock *, 4> NewEndTryBBs;
@@ -2372,6 +2395,48 @@ static void appendEndToFunction(MachineFunction &MF,
           TII.get(WebAssembly::END_FUNCTION));
 }
 
+// We added block~end_block and try_table~end_try_table markers in
+// placeTryTableMarker. But When catch clause's destination has a return type,
+// as in the case of catch with a concrete tag, catch_ref, and catch_all_ref.
+// For example:
+// block exnref
+//   try_table (catch_all_ref 0)
+//     ...
+//   end_try_table
+// end_block
+// ... use exnref ...
+//
+// This code is not valid because the block's body type is not exnref. So we add
+// an unreachable after the 'end_try_table' to make the code valid here:
+// block exnref
+//   try_table (catch_all_ref 0)
+//     ...
+//   end_try_table
+//   unreachable      (new)
+// end_block
+//
+// Because 'unreachable' is a terminator we also need to split the BB.
+static void addUnreachableAfterTryTables(MachineFunction &MF,
+                                         const WebAssemblyInstrInfo &TII) {
+  std::vector<MachineInstr *> EndTryTables;
+  for (auto &MBB : MF)
+    for (auto &MI : MBB)
+      if (MI.getOpcode() == WebAssembly::END_TRY_TABLE)
+        EndTryTables.push_back(&MI);
+
+  for (auto *EndTryTable : EndTryTables) {
+    auto *MBB = EndTryTable->getParent();
+    auto *NewEndTryTableBB = MF.CreateMachineBasicBlock();
+    MF.insert(MBB->getIterator(), NewEndTryTableBB);
+    auto SplitPos = std::next(EndTryTable->getIterator());
+    NewEndTryTableBB->splice(NewEndTryTableBB->end(), MBB, MBB->begin(),
+                             SplitPos);
+    NewEndTryTableBB->addSuccessor(MBB);
+    BuildMI(NewEndTryTableBB, EndTryTable->getDebugLoc(),
+            TII.get(WebAssembly::UNREACHABLE));
+  }
+}
+
 /// Insert BLOCK/LOOP/TRY/TRY_TABLE markers at appropriate places.
 void WebAssemblyCFGStackify::placeMarkers(MachineFunction &MF) {
   // We allocate one more than the number of blocks in the function to
@@ -2398,13 +2463,17 @@ void WebAssemblyCFGStackify::placeMarkers(MachineFunction &MF) {
     }
   }
 
-  // Fix mismatches in unwind destinations induced by linearizing the code.
   if (MCAI->getExceptionHandlingType() == ExceptionHandling::Wasm &&
       MF.getFunction().hasPersonalityFn()) {
-    bool MismatchFixed = fixCallUnwindMismatches(MF);
-    MismatchFixed |= fixCatchUnwindMismatches(MF);
-    if (MismatchFixed)
-      recalculateScopeTops(MF);
+    const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
+    // Add an 'unreachable' after 'end_try_table's.
+    addUnreachableAfterTryTables(MF, TII);
+    // Fix mismatches in unwind destinations induced by linearizing the code.
+    fixCallUnwindMismatches(MF);
+    fixCatchUnwindMismatches(MF);
+    // addUnreachableAfterTryTables and fixUnwindMismatches create new BBs, so
+    // we need to recalculate ScopeTops.
+    recalculateScopeTops(MF);
   }
 }