[WebAssembly] Fix fixEndsAtEndOfFunction for try-catch

[aheejin]

Summary

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 catches, 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.

[aheejin]

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

1105–1109

This comment block is out-of-date; we don't search for 'local.set' or 'drop' anymore. Deleted it while touching this code.

1110–1121

This is extracted as findCatch method above.

[aheejin]

• clang-tidy fix

[tlively]

LGTM!

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

 #endif
       break;
     }
     ++InsertPos;
   }
   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;
   EndToBegin[End] = Begin;
 }
 
 void WebAssemblyCFGStackify::registerTryScope(MachineInstr *Begin,
                                               MachineInstr *End,
for (auto &P : UnwindDestToTryRanges) {
     NumUnwindMismatches += P.second.size();
 
     // This means the destination is the appendix BB, which was separately
     // handled above.
     if (!P.first)
       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
     // inner mismatched scope.
     MachineInstr *BeginTry = EHPadToTry[EHPad];
     MachineInstr *EndTry = BeginToEnd[BeginTry];
     MachineBasicBlock *Cont = EndTry->getParent();
     auto *BrDest = MF.CreateMachineBasicBlock();
void WebAssemblyCFGStackify::fixEndsAtEndOfFunction(MachineFunction &MF) {
   // MCInstLower will add the proper types to multivalue signatures based on the
   // function return type
   WebAssembly::BlockType RetType =
       MFI.getResults().size() > 1
           ? WebAssembly::BlockType::Multivalue
           : 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
 // were a block.
 static void appendEndToFunction(MachineFunction &MF,
                                 const WebAssemblyInstrInfo &TII) {
   BuildMI(MF.back(), MF.back().end(),
           MF.back().findPrevDebugLoc(MF.back().end()),

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

ehcleanup:                                        ; preds = %catch.start
   %6 = call i8* @llvm.wasm.get.exception(token %5)
   call void @__clang_call_terminate(i8* %6) [ "funclet"(token %5) ]
   unreachable
 
 while.end:                                        ; preds = %while.body, %while.cond
   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
 
 declare void @foo()
 declare void @bar()
 declare i32 @baz()
 declare i32 @qux(i32)
 declare void @quux(i32)