[WebAssembly] Fix remapping branch dests in fixCatchUnwindMismatches

[aheejin]

Summary

This is a case D97178 tried to solve but missed. D97178 could not handle
the case when
multiple consecutive delegates are generated:

• Before:

block
  br (a)
  try
  catch
  end_try
end_block
          <- (a)

• After

block
  br (a)
  try
    ...
    try
      try
      catch
      end_try
            <- (a)
    delegate
  delegate
end_block
          <- (b)

(The br should point to (b) now)

D97178 assumed end_block exists two BBs later than end_try, because
it assumed the order as end_try BB -> delegate BB -> end_block BB.
But it turned out there can be multiple delegates in between. This
patch changes the logic so we just search from end_try BB until we
find end_block.

Fixes https://github.com/emscripten-core/emscripten/issues/13515.
(More precisely, fixes
https://github.com/emscripten-core/emscripten/issues/13515#issuecomment-784711318.)

[aheejin]

whitespace fix

[tlively]

Makes sense!

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

1200

genearted => generated

1223

destnation => destination

[aheejin]

@tlively Thanks! bugpoint really helped reducing this case.

[aheejin]

Address comments

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

for (auto &MI : reverse(MBB)) {
       }
     }
   }
 
   assert(EHPadStack.empty());
   if (EHPadToUnwindDest.empty())
     return false;
   NumCatchUnwindMismatches += EHPadToUnwindDest.size();
-  // <current branch dest, future branch dest> map, because fixing catch unwind
-  // mismatches can invalidate branch destinations
-  DenseMap<MachineBasicBlock *, MachineBasicBlock *> BrDestMap;
+  SmallPtrSet<MachineBasicBlock *, 4> NewEndTryBBs;
 
   for (auto &P : EHPadToUnwindDest) {
     MachineBasicBlock *EHPad = P.first;
     MachineBasicBlock *UnwindDest = P.second;
     MachineInstr *Try = EHPadToTry[EHPad];
     MachineInstr *EndTry = BeginToEnd[Try];
     addTryDelegate(Try, EndTry, UnwindDest);
-    BrDestMap[EndTry->getParent()] =
-        EndTry->getParent()->getNextNode()->getNextNode();
+    NewEndTryBBs.insert(EndTry->getParent());
   }
 
   // Adding a try-delegate wrapping an existing try-catch-end can make existing
   // branch destination BBs invalid. For example,
   //
   // - Before:
   // bb0:
   //   block
bool WebAssemblyCFGStackify::fixCatchUnwindMismatches(MachineFunction &MF) {
   // split_bb:     ;; (new BB)
   //   end_block
   //
   // Now 'br bb3' incorrectly branches to an inner scope.
   //
   // As we can see in this case, when branches target a BB that has both
   // 'end_try' and 'end_block' and the BB is split to insert a 'delegate', we
   // have to remap existing branch destinations so that they target not the
-  // 'end_try' BB but the new 'end_block' BB, which should be the second next BB
-  // of 'end_try' (because there is a 'delegate' BB in between). In this
-  // example, the 'br bb3' instruction should be remapped to 'br split_bb'.
+  // 'end_try' BB but the new 'end_block' BB. There can be multiple 'delegate's
+  // in between, so we try to find the next BB with 'end_block' instruction. In
+  // this example, the 'br bb3' instruction should be remapped to 'br split_bb'.
   for (auto &MBB : MF) {
     for (auto &MI : MBB) {
       if (MI.isTerminator()) {
         for (auto &MO : MI.operands()) {
-          if (MO.isMBB()) {
-            auto It = BrDestMap.find(MO.getMBB());
-            if (It != BrDestMap.end())
-              MO.setMBB(It->second);
+          if (MO.isMBB() && NewEndTryBBs.count(MO.getMBB())) {
+            auto *BrDest = MO.getMBB();
+            bool FoundEndBlock = false;
+            for (; std::next(BrDest->getIterator()) != MF.end();
+                 BrDest = BrDest->getNextNode()) {
+              for (const auto &MI : *BrDest) {
+                if (MI.getOpcode() == WebAssembly::END_BLOCK) {
+                  FoundEndBlock = true;
+                  break;
+                }
+              }
+              if (FoundEndBlock)
+                break;
+            }
+            assert(FoundEndBlock);
+            MO.setMBB(BrDest);
           }
         }
       }
     }
   }
 
   return true;
 }

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

catch.start1:                                     ; preds = %catch.dispatch1
   %6 = call i8* @llvm.wasm.get.exception(token %5)
   %7 = call i32 @llvm.wasm.get.ehselector(token %5)
   catchret from %5 to label %try.cont
 
 try.cont:                                         ; preds = %catch.start1, %catch.start0, %bb1
   ret void
 }
 
+; The similar case with test20, but multiple consecutive delegates are
+; generated:
+; - Before:
+; block
+;   br (a)
+;   try
+;   catch
+;   end_try
+; end_block
+;           <- (a)
+;
+; - After
+; block
+;   br (a)
+;   try
+;     ...
+;     try
+;       try
+;       catch
+;       end_try
+;             <- (a)
+;     delegate
+;   delegate
+; end_block
+;           <- (b) The br destination should be remapped to here
+;
+; The test was reduced by bugpoint and should not crash.
+define void @test21() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+entry:
+  br i1 undef, label %if.then, label %if.end12
+
+if.then:                                          ; preds = %entry
+  invoke void @__cxa_throw() #1
+          to label %unreachable unwind label %catch.dispatch
+
+catch.dispatch:                                   ; preds = %if.then
+  %0 = catchswitch within none [label %catch.start] unwind to caller
+
+catch.start:                                      ; preds = %catch.dispatch
+  %1 = catchpad within %0 [i8* bitcast (i8** @_ZTIi to i8*)]
+  %2 = call i8* @llvm.wasm.get.exception(token %1)
+  %3 = call i32 @llvm.wasm.get.ehselector(token %1)
+  catchret from %1 to label %catchret.dest
+
+catchret.dest:                                    ; preds = %catch.start
+  invoke void @foo()
+          to label %invoke.cont unwind label %catch.dispatch4
+
+invoke.cont:                                      ; preds = %catchret.dest
+  invoke void @__cxa_throw() #1
+          to label %unreachable unwind label %catch.dispatch4
+
+catch.dispatch4:                                  ; preds = %invoke.cont, %catchret.dest
+  %4 = catchswitch within none [label %catch.start5] unwind to caller
+
+catch.start5:                                     ; preds = %catch.dispatch4
+  %5 = catchpad within %4 [i8* bitcast (i8** @_ZTIi to i8*)]
+  %6 = call i8* @llvm.wasm.get.exception(token %5)
+  %7 = call i32 @llvm.wasm.get.ehselector(token %5)
+  unreachable
+
+if.end12:                                         ; preds = %entry
+  invoke void @foo()
+          to label %invoke.cont14 unwind label %catch.dispatch16
+
+catch.dispatch16:                                 ; preds = %if.end12
+  %8 = catchswitch within none [label %catch.start17] unwind label %ehcleanup
+
+catch.start17:                                    ; preds = %catch.dispatch16
+  %9 = catchpad within %8 [i8* bitcast (i8** @_ZTIi to i8*)]
+  %10 = call i8* @llvm.wasm.get.exception(token %9)
+  %11 = call i32 @llvm.wasm.get.ehselector(token %9)
+  br i1 undef, label %catch20, label %rethrow19
+
+catch20:                                          ; preds = %catch.start17
+  catchret from %9 to label %catchret.dest22
+
+catchret.dest22:                                  ; preds = %catch20
+  br label %try.cont23
+
+rethrow19:                                        ; preds = %catch.start17
+  invoke void @llvm.wasm.rethrow() #1 [ "funclet"(token %9) ]
+          to label %unreachable unwind label %ehcleanup
+
+try.cont23:                                       ; preds = %invoke.cont14, %catchret.dest22
+  invoke void @foo()
+          to label %invoke.cont24 unwind label %ehcleanup
+
+invoke.cont24:                                    ; preds = %try.cont23
+  ret void
+
+invoke.cont14:                                    ; preds = %if.end12
+  br label %try.cont23
+
+ehcleanup:                                        ; preds = %try.cont23, %rethrow19, %catch.dispatch16
+  %12 = cleanuppad within none []
+  cleanupret from %12 unwind to caller
+
+unreachable:                                      ; preds = %if.then, %invoke.cont, %rethrow19
+  unreachable
+}
+
 ; Check if the unwind destination mismatch stats are correct
 ; NOSORT: 20 wasm-cfg-stackify    - Number of call unwind mismatches found
 ; NOSORT:  4 wasm-cfg-stackify    - Number of catch unwind mismatches found
 
 declare void @foo()
 declare void @bar()
 declare i32 @baz()
 declare i32 @qux(i32)

 declare %class.MyClass* @_ZN7MyClassC2ERKS_(%class.MyClass* returned, %class.MyClass* dereferenceable(4)) #0
 
 declare i32 @__gxx_wasm_personality_v0(...)
 ; Function Attrs: nounwind
 declare i8* @llvm.wasm.get.exception(token) #0
 ; Function Attrs: nounwind
 declare i32 @llvm.wasm.get.ehselector(token) #0
 ; Function Attrs: noreturn
+declare void @__cxa_throw() #1
+; Function Attrs: noreturn
 declare void @llvm.wasm.rethrow() #1
 ; Function Attrs: nounwind
 declare i32 @llvm.eh.typeid.for(i8*) #0
 
 declare i8* @__cxa_begin_catch(i8*)
 declare void @__cxa_end_catch()
 declare i8* @__cxa_get_exception_ptr(i8*)
 declare void @__clang_call_terminate(i8*)