[WebAssembly] Fix a bug in finding matching EH pad

[aheejin]

Summary

getMatchingEHPad() in LateEHPrepare is a function to find the nearest
EH pad that dominates the given instruction. This intends to be
lightweight so it does not use full WebAssemblyException scope analysis
or dominator analysis. It simply does backward BFS to its predecessors
and stops at the first EH pad each search path encounters. All search
should end up at the same EH pad, and if not, it returns null.

But it didn't take into account that when there are inner scopes within
the current scope, some path in BFS can hit an inner EH pad first. For
example, in the given diagram, Inst belongs to the outer scope and
getMathingEHPad() should return 'EHPad 1', but some search path can go
into the inner scope and end up with 'EHPad 2'. The search will return
null because different paths end up with different EH pads.

--- EHPad 1 ---
| - EHPad 2 - |
| |         | |
| ----------- |
|   Inst      |
---------------

So far this was OK because we haven't tested a case in which a given
instruction is far from its EH pad. Also, this bug does not happen when
the inner EH scope is a cleanup scope, because a cleanup scope ends with
a cleanupret whose successor is an EH pad, so the search encounters
that EH pad first before going into the child scope. But this can happen
when the child scope is a catch scope that ends with catchret. So this
patch, when doing backward BFS, does not search predecessors that ends
with catchret. Because catchrets are replaced with brs during this
pass, this records BBs that have catchrets in the beginning, before
doing any other transformations.

[aheejin]

• Fix test function name

llvm/lib/Target/WebAssembly/WebAssemblyLateEHPrepare.cpp

 
 namespace {
 class WebAssemblyLateEHPrepare final : public MachineFunctionPass {
   StringRef getPassName() const override {
     return "WebAssembly Late Prepare Exception";
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
+  void recordCatchRetBBs(MachineFunction &MF);
   bool addCatches(MachineFunction &MF);
   bool replaceFuncletReturns(MachineFunction &MF);
   bool removeUnnecessaryUnreachables(MachineFunction &MF);
   bool addExceptionExtraction(MachineFunction &MF);
   bool restoreStackPointer(MachineFunction &MF);
 
+  MachineBasicBlock *getMatchingEHPad(MachineInstr *MI);
+  SmallSet<MachineBasicBlock *, 8> CatchRetBBs;
+
 public:
   static char ID; // Pass identification, replacement for typeid
   WebAssemblyLateEHPrepare() : MachineFunctionPass(ID) {}
 };
 } // end anonymous namespace
 
 char WebAssemblyLateEHPrepare::ID = 0;
 INITIALIZE_PASS(WebAssemblyLateEHPrepare, DEBUG_TYPE,
                 "WebAssembly Late Exception Preparation", false, false)
 
 FunctionPass *llvm::createWebAssemblyLateEHPrepare() {
   return new WebAssemblyLateEHPrepare();
 }
 
 // Returns the nearest EH pad that dominates this instruction. This does not use
 // dominator analysis; it just does BFS on its predecessors until arriving at an
 // EH pad. This assumes valid EH scopes so the first EH pad it arrives in all
 // possible search paths should be the same.
 // Returns nullptr in case it does not find any EH pad in the search, or finds
 // multiple different EH pads.
-static MachineBasicBlock *getMatchingEHPad(MachineInstr *MI) {
+MachineBasicBlock *
+WebAssemblyLateEHPrepare::getMatchingEHPad(MachineInstr *MI) {
   MachineFunction *MF = MI->getParent()->getParent();
   SmallVector<MachineBasicBlock *, 2> WL;
   SmallPtrSet<MachineBasicBlock *, 2> Visited;
   WL.push_back(MI->getParent());
   MachineBasicBlock *EHPad = nullptr;
   while (!WL.empty()) {
     MachineBasicBlock *MBB = WL.pop_back_val();
     if (Visited.count(MBB))
       continue;
     Visited.insert(MBB);
     if (MBB->isEHPad()) {
       if (EHPad && EHPad != MBB)
         return nullptr;
       EHPad = MBB;
       continue;
     }
     if (MBB == &MF->front())
       return nullptr;
-    WL.append(MBB->pred_begin(), MBB->pred_end());
+    for (auto *Pred : MBB->predecessors())
+      if (!CatchRetBBs.count(Pred)) // We don't go into child scopes
+        WL.push_back(Pred);
   }
   return EHPad;
 }
 
 // Erase the specified BBs if the BB does not have any remaining predecessors,
 // and also all its dead children.
 template <typename Container>
 static void eraseDeadBBsAndChildren(const Container &MBBs) {
LLVM_DEBUG(dbgs() << "********** Late EH Prepare **********\n"
                     << MF.getName() << '\n');
 
   if (MF.getTarget().getMCAsmInfo()->getExceptionHandlingType() !=
       ExceptionHandling::Wasm)
     return false;
 
   bool Changed = false;
   if (MF.getFunction().hasPersonalityFn()) {
+    recordCatchRetBBs(MF);
     Changed |= addCatches(MF);
     Changed |= replaceFuncletReturns(MF);
   }
   Changed |= removeUnnecessaryUnreachables(MF);
   if (MF.getFunction().hasPersonalityFn()) {
     Changed |= addExceptionExtraction(MF);
     Changed |= restoreStackPointer(MF);
   }
   return Changed;
 }
 
+// Record which BB ends with 'CATCHRET' instruction, because this will be
+// replaced with BRs later. This set of 'CATCHRET' BBs is necessary in
+// 'getMatchingEHPad' function.
+void WebAssemblyLateEHPrepare::recordCatchRetBBs(MachineFunction &MF) {
+  CatchRetBBs.clear();
+  for (auto &MBB : MF) {
+    auto Pos = MBB.getFirstTerminator();
+    if (Pos == MBB.end())
+      continue;
+    MachineInstr *TI = &*Pos;
+    if (TI->getOpcode() == WebAssembly::CATCHRET)
+      CatchRetBBs.insert(&MBB);
+  }
+}
+
 // Add catch instruction to beginning of catchpads and cleanuppads.
 bool WebAssemblyLateEHPrepare::addCatches(MachineFunction &MF) {
   bool Changed = false;
   const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
   MachineRegisterInfo &MRI = MF.getRegInfo();
   for (auto &MBB : MF) {
     if (MBB.isEHPad()) {
       Changed = true;

llvm/test/CodeGen/WebAssembly/exception.ll

catch:                                            ; preds = %catch.start
   %5 = call i8* @__cxa_begin_catch(i8* %2) [ "funclet"(token %1) ]
   call void @__cxa_end_catch() [ "funclet"(token %1) ]
   catchret from %1 to label %try.cont
 
 rethrow:                                          ; preds = %catch.start
   call void @llvm.wasm.rethrow.in.catch() [ "funclet"(token %1) ]
   unreachable
 
-try.cont:                                         ; preds = %entry, %catch
+try.cont:                                         ; preds = %catch, %entry
   ret void
 }
 
 ; Destructor (cleanup) test
 ;
 ; void foo();
 ; struct Temp {
 ;   ~Temp() {}
catch.start:                                      ; preds = %catch.dispatch
   %4 = call i8* @__cxa_begin_catch(i8* %2) [ "funclet"(token %1) ]
   invoke void @foo() [ "funclet"(token %1) ]
           to label %invoke.cont1 unwind label %ehcleanup
 
 invoke.cont1:                                     ; preds = %catch.start
   call void @__cxa_end_catch() [ "funclet"(token %1) ]
   catchret from %1 to label %try.cont
 
-try.cont:                                         ; preds = %entry, %invoke.cont1
+try.cont:                                         ; preds = %invoke.cont1, %entry
   ret void
 
 ehcleanup:                                        ; preds = %catch.start
   %5 = cleanuppad within %1 []
   invoke void @__cxa_end_catch() [ "funclet"(token %5) ]
           to label %invoke.cont2 unwind label %terminate
 
 invoke.cont2:                                     ; preds = %ehcleanup

 invoke.cont1:                                     ; preds = %catch
   call void @__cxa_end_catch() [ "funclet"(token %1) ]
   catchret from %1 to label %try.cont
 
 rethrow:                                          ; preds = %catch.start
   call void @llvm.wasm.rethrow.in.catch() [ "funclet"(token %1) ]
   unreachable
 
-try.cont:                                         ; preds = %entry, %invoke.cont1
+try.cont:                                         ; preds = %invoke.cont1, %entry
   ret void
 
 ehcleanup:                                        ; preds = %catch
   %8 = cleanuppad within %1 []
   call void @__cxa_end_catch() [ "funclet"(token %8) ]
   cleanupret from %8 unwind to caller
 }
 

 
 catch.dispatch:                                   ; preds = %entry
   %0 = catchswitch within none [label %catch.start] unwind to caller
 
 catch.start:                                      ; preds = %catch.dispatch
   %1 = catchpad within %0 [i8* null]
   %2 = call i8* @llvm.wasm.get.exception(token %1)
   %3 = call i32 @llvm.wasm.get.ehselector(token %1)
-  %4 = call i8* @__cxa_begin_catch(i8* %2) #2 [ "funclet"(token %1) ]
+  %4 = call i8* @__cxa_begin_catch(i8* %2) [ "funclet"(token %1) ]
   call void @__cxa_end_catch() [ "funclet"(token %1) ]
   catchret from %1 to label %try.cont
 
-try.cont:                                         ; preds = %entry, %catch.start
+try.cont:                                         ; preds = %catch.start, %entry
   ret void
 }
 
 ; When the result of @llvm.wasm.get.exception is not used. This is created to
 ; fix a bug in LateEHPrepare and this should not crash.
 define void @test_get_exception_wo_use() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 entry:
   invoke void @foo()
           to label %try.cont unwind label %catch.dispatch
 
 catch.dispatch:                                   ; preds = %entry
   %0 = catchswitch within none [label %catch.start] unwind to caller
 
 catch.start:                                      ; preds = %catch.dispatch
   %1 = catchpad within %0 [i8* null]
   %2 = call i8* @llvm.wasm.get.exception(token %1)
   %3 = call i32 @llvm.wasm.get.ehselector(token %1)
   catchret from %1 to label %try.cont
 
-try.cont:                                         ; preds = %entry, %catch.start
+try.cont:                                         ; preds = %catch.start, %entry
+  ret void
+}
+
+; Tests a case when a cleanup region (cleanuppad ~ clanupret) contains another
+; catchpad
+define void @test_complex_cleanup_region() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+entry:
+  invoke void @foo()
+          to label %invoke.cont unwind label %ehcleanup
+
+invoke.cont:                                      ; preds = %entry
   ret void
+
+ehcleanup:                                        ; preds = %entry
+  %0 = cleanuppad within none []
+  invoke void @foo() [ "funclet"(token %0) ]
+          to label %ehcleanupret unwind label %catch.dispatch
+
+catch.dispatch:                                   ; preds = %ehcleanup
+  %1 = catchswitch within %0 [label %catch.start] unwind label %ehcleanup.1
+
+catch.start:                                      ; preds = %catch.dispatch
+  %2 = catchpad within %1 [i8* null]
+  %3 = call i8* @llvm.wasm.get.exception(token %2)
+  %4 = call i32 @llvm.wasm.get.ehselector(token %2)
+  catchret from %2 to label %ehcleanupret
+
+ehcleanup.1:                                      ; preds = %catch.dispatch
+  %5 = cleanuppad within %0 []
+  unreachable
+
+ehcleanupret:                                     ; preds = %catch.start, %ehcleanup
+  cleanupret from %0 unwind to caller
 }
 
 declare void @foo()
 declare void @bar(i32*)
 declare i32 @__gxx_wasm_personality_v0(...)
 declare void @llvm.wasm.throw(i32, i8*)
 declare i8* @llvm.wasm.get.exception(token)
 declare i32 @llvm.wasm.get.ehselector(token)