[WebAssemblly] Fix rethrow's argument computation

[aheejin]

Summary

Previously we assumed rethrow's argument was always 0, but it turned
out rethrow follows the same rule with br or delegate:
https://github.com/WebAssembly/exception-handling/pull/137
https://github.com/WebAssembly/exception-handling/issues/146#issuecomment-777349038

Currently rethrows generated by our backend always rethrow the
exception caught by the innermost enclosing catch, so this adds a
function to compute that and replaces rethrow's argument with its
computed result.

[aheejin]

Sorry, this had a bug and I had to change the code significantly. To prevent confusion I'll close this and open a new one.

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

class WebAssemblyCFGStackify final : public MachineFunctionPass {
 
   // Wrap-up
   using EndMarkerInfo =
       std::pair<const MachineBasicBlock *, const MachineInstr *>;
   unsigned getBranchDepth(const SmallVectorImpl<EndMarkerInfo> &Stack,
                           const MachineBasicBlock *MBB);
   unsigned getDelegateDepth(const SmallVectorImpl<EndMarkerInfo> &Stack,
                             const MachineBasicBlock *MBB);
+  unsigned getRethrowDepth(const SmallVectorImpl<EndMarkerInfo> &Stack);
   void rewriteDepthImmediates(MachineFunction &MF);
   void fixEndsAtEndOfFunction(MachineFunction &MF);
   void cleanupFunctionData(MachineFunction &MF);
 
   // For each BLOCK|LOOP|TRY, the corresponding END_(BLOCK|LOOP|TRY) or DELEGATE
   // (in case of TRY).
   DenseMap<const MachineInstr *, MachineInstr *> BeginToEnd;
   // For each END_(BLOCK|LOOP|TRY) or DELEGATE, the corresponding
for (auto X : reverse(Stack)) {
     if (X.first == EndTry->getParent() && X.second == EndTry)
       break;
     ++Depth;
   }
   assert(Depth < Stack.size() && "Delegate destination should be in scope");
   return Depth;
 }
 
+unsigned WebAssemblyCFGStackify::getRethrowDepth(
+    const SmallVectorImpl<EndMarkerInfo> &Stack) {
+  unsigned Depth = 0;
+  // In our current implementation, rethrows always rethrow the exception caught
+  // by the innermost enclosing catch. So find the nearest end_try marker.
+  for (auto X : reverse(Stack)) {
+    if (X.second->getOpcode() == WebAssembly::END_TRY)
+      break;
+    ++Depth;
+  }
+  assert(Depth < Stack.size() && "Rethrow destination should be in scope");
+  return Depth;
+}
+
 void WebAssemblyCFGStackify::rewriteDepthImmediates(MachineFunction &MF) {
   // Now rewrite references to basic blocks to be depth immediates.
   SmallVector<EndMarkerInfo, 8> Stack;
   for (auto &MBB : reverse(MF)) {
     for (auto I = MBB.rbegin(), E = MBB.rend(); I != E; ++I) {
       MachineInstr &MI = *I;
       switch (MI.getOpcode()) {
       case WebAssembly::BLOCK:
for (auto I = MBB.rbegin(), E = MBB.rend(); I != E; ++I) {
         // immediate operands to rewrite.
         Stack.push_back(std::make_pair(&MBB, &MI));
         break;
 
       case WebAssembly::END_LOOP:
         Stack.push_back(std::make_pair(EndToBegin[&MI]->getParent(), &MI));
         break;
 
+      case WebAssembly::RETHROW:
+        MI.getOperand(0).setImm(getRethrowDepth(Stack));
+        break;
+
       default:
         if (MI.isTerminator()) {
           // Rewrite MBB operands to be depth immediates.
           SmallVector<MachineOperand, 4> Ops(MI.operands());
           while (MI.getNumOperands() > 0)
             MI.RemoveOperand(MI.getNumOperands() - 1);
           for (auto MO : Ops) {
             if (MO.isMBB()) {

llvm/lib/Target/WebAssembly/WebAssemblyInstrControl.td

 // Throwing an exception: throw / rethrow
 let isTerminator = 1, hasCtrlDep = 1, isBarrier = 1 in {
 defm THROW : I<(outs), (ins event_op:$tag, variable_ops),
                (outs), (ins event_op:$tag),
                [(WebAssemblythrow (WebAssemblywrapper texternalsym:$tag))],
                "throw   \t$tag", "throw   \t$tag", 0x08>;
 defm RETHROW : NRI<(outs), (ins i32imm:$depth), [], "rethrow \t$depth", 0x09>;
 } // isTerminator = 1, hasCtrlDep = 1, isBarrier = 1
-// For C++ support, we only rethrow the latest exception, thus always setting
-// the depth to 0.
+// The depth argument will be computed in CFGStackify. We set it to 0 here for
+// now.
 def : Pat<(int_wasm_rethrow), (RETHROW 0)>;
 
 // Region within which an exception is caught: try / end_try
 let Uses = [VALUE_STACK], Defs = [VALUE_STACK] in {
 defm TRY     : NRI<(outs), (ins Signature:$sig), [], "try     \t$sig", 0x06>;
 defm END_TRY : NRI<(outs), (ins), [], "end_try", 0x0b>;
 } // Uses = [VALUE_STACK], Defs = [VALUE_STACK]
 

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

 ; CHECK:         catch_all                             # catch[[C0]]:
 ; CHECK:           call      __cxa_end_catch
 ; CHECK:           rethrow   0                         # to caller
 ; CHECK:         end_try                               # label[[L3]]:
 ; CHECK:         call      __cxa_end_catch
 ; CHECK:         br        2                           # 2: down to label[[L1]]
 ; CHECK:       end_try
 ; CHECK:     end_block                                 # label[[L0]]:
-; CHECK:     rethrow   0                               # to caller
+; CHECK:     rethrow   1                               # to caller
 ; CHECK:   end_block                                   # label[[L1]]:
 ; CHECK:   call      __cxa_end_catch
 ; CHECK: end_try
 define void @test1() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 entry:
   invoke void @foo()
           to label %try.cont11 unwind label %catch.dispatch