[WebAssembly] Fix fixEndsAtEndOfFunction for try-catch

When the function return type is non-void and `end` instructions are at
the very end of a function, CFGStackify's `fixEndsAtEndOfFunction`
function fixes the corresponding block/loop/try's type to match the
function's return type. This is applied to consecutive `end` markers at
the end of a function. For example, when the function return type is
`i32`,
```
block i32    ;; return type is fixed to i32
  ...
  loop i32   ;; return type is fixed to i32
    ...
  end_loop
end_block
end_function
```

But try-catch is a little different, because it consists of two parts:
a try part and a catch part, and both parts' return type should satisfy
the function's return type. Which means,
```
try i32      ;; return type is fixed to i32
  ...
  block i32  ;; this should be changed i32 too!
    ...
  end_block
catch
  ...
end_try
end_function
```
As you can see in this example, it is not sufficient to only `end`
instructions at the end of a function; in case of `try`, we should
check instructions before `catch`es, in case their corresponding `try`'s
type has been fixed.

This changes `fixEndsAtEndOfFunction`'s algorithm to use a worklist
that contains a reverse iterator, each of which is a starting point for
a new backward `end` instruction search.

Fixes https://bugs.llvm.org/show_bug.cgi?id=47413.

Reviewed By: dschuff, tlively

Differential Revision: https://reviews.llvm.org/D87207

1 files changed, 48 insertions, 24 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
index 02330a2..d5ee4b3 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
@@ -178,6 +178,28 @@ getLatestInsertPos(MachineBasicBlock *MBB,
   return InsertPos;
 }
 
+// Find a catch instruction and its destination register within an EH pad.
+static MachineInstr *findCatch(MachineBasicBlock *EHPad, Register &ExnReg) {
+  assert(EHPad->isEHPad());
+  MachineInstr *Catch = nullptr;
+  for (auto &MI : *EHPad) {
+    switch (MI.getOpcode()) {
+    case WebAssembly::CATCH:
+      Catch = &MI;
+      ExnReg = Catch->getOperand(0).getReg();
+      break;
+    }
+  }
+  assert(Catch && "EH pad does not have a catch");
+  assert(ExnReg != 0 && "Invalid register");
+  return Catch;
+}
+
+static MachineInstr *findCatch(MachineBasicBlock *EHPad) {
+  Register Dummy;
+  return findCatch(EHPad, Dummy);
+}
+
 void WebAssemblyCFGStackify::registerScope(MachineInstr *Begin,
                                            MachineInstr *End) {
   BeginToEnd[Begin] = End;
@@ -1101,25 +1123,8 @@ bool WebAssemblyCFGStackify::fixUnwindMismatches(MachineFunction &MF) {
       continue;
 
     MachineBasicBlock *EHPad = P.first;
-
-    // Find 'catch' and 'local.set' or 'drop' instruction that follows the
-    // 'catch'. If -wasm-disable-explicit-locals is not set, 'catch' should be
-    // always followed by either 'local.set' or a 'drop', because 'br_on_exn' is
-    // generated after 'catch' in LateEHPrepare and we don't support blocks
-    // taking values yet.
-    MachineInstr *Catch = nullptr;
-    unsigned ExnReg = 0;
-    for (auto &MI : *EHPad) {
-      switch (MI.getOpcode()) {
-      case WebAssembly::CATCH:
-        Catch = &MI;
-        ExnReg = Catch->getOperand(0).getReg();
-        break;
-      }
-    }
-    assert(Catch && "EH pad does not have a catch");
-    assert(ExnReg != 0 && "Invalid register");
-
+    Register ExnReg = 0;
+    MachineInstr *Catch = findCatch(EHPad, ExnReg);
     auto SplitPos = std::next(Catch->getIterator());
 
     // Create a new BB that's gonna be the destination for branches from the
@@ -1371,22 +1376,41 @@ void WebAssemblyCFGStackify::fixEndsAtEndOfFunction(MachineFunction &MF) {
           : WebAssembly::BlockType(
                 WebAssembly::toValType(MFI.getResults().front()));
 
-  for (MachineBasicBlock &MBB : reverse(MF)) {
-    for (MachineInstr &MI : reverse(MBB)) {
+  SmallVector<MachineBasicBlock::reverse_iterator, 4> Worklist;
+  Worklist.push_back(MF.rbegin()->rbegin());
+
+  auto Process = [&](MachineBasicBlock::reverse_iterator It) {
+    auto *MBB = It->getParent();
+    while (It != MBB->rend()) {
+      MachineInstr &MI = *It++;
       if (MI.isPosition() || MI.isDebugInstr())
         continue;
       switch (MI.getOpcode()) {
+      case WebAssembly::END_TRY: {
+        // If a 'try''s return type is fixed, both its try body and catch body
+        // should satisfy the return type, so we need to search 'end'
+        // instructions before its corresponding 'catch' too.
+        auto *EHPad = TryToEHPad.lookup(EndToBegin[&MI]);
+        assert(EHPad);
+        Worklist.push_back(std::next(findCatch(EHPad)->getReverseIterator()));
+        LLVM_FALLTHROUGH;
+      }
       case WebAssembly::END_BLOCK:
       case WebAssembly::END_LOOP:
-      case WebAssembly::END_TRY:
         EndToBegin[&MI]->getOperand(0).setImm(int32_t(RetType));
         continue;
       default:
-        // Something other than an `end`. We're done.
+        // Something other than an `end`. We're done for this BB.
         return;
       }
     }
-  }
+    // We've reached the beginning of a BB. Continue the search in the previous
+    // BB.
+    Worklist.push_back(MBB->getPrevNode()->rbegin());
+  };
+
+  while (!Worklist.empty())
+    Process(Worklist.pop_back_val());
 }
 
 // WebAssembly functions end with an end instruction, as if the function body