[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

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

test/CodeGen/WebAssembly/cfg-stackify-eh.ll

@@ -1023,6 +1023,54 @@ while.end:                                        ; preds = %while.body, %while.
   ret void
 }
 
+; 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.
+; But when a `try`'s type is fixed, we should also check `end` instructions
+; before its corresponding `catch`, because both `try` and `catch` body should
+; satisfy the return type requirements.
+
+; NOSORT-LABEL: test19
+; NOSORT: try i32
+; NOSORT: loop i32
+; NOSORT: end_loop
+; NOSORT: catch
+; NOSORT: end_try
+; NOSORT-NEXT: end_function
+define i32 @test19(i32 %n) personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+entry:
+  %t = alloca %class.Object, align 1
+  br label %for.cond
+
+for.cond:                                         ; preds = %for.inc, %entry
+  %i.0 = phi i32 [ 0, %entry ], [ %inc, %for.inc ]
+  %cmp = icmp slt i32 %i.0, %n
+  br label %for.body
+
+for.body:                                         ; preds = %for.cond
+  %div = sdiv i32 %n, 2
+  %cmp1 = icmp eq i32 %i.0, %div
+  br i1 %cmp1, label %if.then, label %for.inc
+
+if.then:                                          ; preds = %for.body
+  %call = invoke i32 @baz()
+          to label %invoke.cont unwind label %ehcleanup
+
+invoke.cont:                                      ; preds = %if.then
+  %call2 = call %class.Object* @_ZN6ObjectD2Ev(%class.Object* %t) #4
+  ret i32 %call
+
+for.inc:                                          ; preds = %for.body
+  %inc = add nsw i32 %i.0, 1
+  br label %for.cond
+
+ehcleanup:                                        ; preds = %if.then
+  %0 = cleanuppad within none []
+  %call3 = call %class.Object* @_ZN6ObjectD2Ev(%class.Object* %t) #4 [ "funclet"(token %0) ]
+  cleanupret from %0 unwind to caller
+}
+
+
 ; Check if the unwind destination mismatch stats are correct
 ; NOSORT-STAT: 17 wasm-cfg-stackify    - Number of EH pad unwind mismatches found