[WebAssemblly] Fix rethrow's argument computation

Previously we assumed `rethrow`'s argument was always 0, but it turned
out `rethrow` follows the same rule with `br` or `delegate`:
WebAssembly/exception-handling#137
WebAssembly/exception-handling#146 (comment)

Currently `rethrow`s 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.

This also renames `EHPadStack` in `InstPrinter` to `TryStack`, because
in CFGStackify we use `EHPadStack` to mean the range between
`catch`~`end`, while in `InstPrinter` we used it to mean the range
between `try`~`catch`, so choosing different names would look clearer.
Doesn't contain any functional changes in `InstPrinter`.

Reviewed By: dschuff

Differential Revision: https://reviews.llvm.org/D96595

llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyInstPrinter.cpp

@@ -104,7 +104,7 @@ void WebAssemblyInstPrinter::printInst(const MCInst *MI, uint64_t Address,
     case WebAssembly::TRY:
     case WebAssembly::TRY_S:
       ControlFlowStack.push_back(std::make_pair(ControlFlowCounter, false));
-      EHPadStack.push_back(ControlFlowCounter++);
+      TryStack.push_back(ControlFlowCounter++);
       EHInstStack.push_back(TRY);
       return;
 
@@ -149,11 +149,10 @@ void WebAssemblyInstPrinter::printInst(const MCInst *MI, uint64_t Address,
       } else if (EHInstStack.back() == CATCH_ALL) {
         printAnnotation(OS, "catch/catch_all cannot occur after catch_all");
       } else if (EHInstStack.back() == TRY) {
-        if (EHPadStack.empty()) {
+        if (TryStack.empty()) {
           printAnnotation(OS, "try-catch mismatch!");
         } else {
-          printAnnotation(OS,
-                          "catch" + utostr(EHPadStack.pop_back_val()) + ':');
+          printAnnotation(OS, "catch" + utostr(TryStack.pop_back_val()) + ':');
         }
         EHInstStack.pop_back();
         if (Opc == WebAssembly::CATCH || Opc == WebAssembly::CATCH_S) {
@@ -168,28 +167,27 @@ void WebAssemblyInstPrinter::printInst(const MCInst *MI, uint64_t Address,
     case WebAssembly::RETHROW_S:
       // 'rethrow' rethrows to the nearest enclosing catch scope, if any. If
       // there's no enclosing catch scope, it throws up to the caller.
-      if (EHPadStack.empty()) {
+      if (TryStack.empty()) {
         printAnnotation(OS, "to caller");
       } else {
-        printAnnotation(OS, "down to catch" + utostr(EHPadStack.back()));
+        printAnnotation(OS, "down to catch" + utostr(TryStack.back()));
       }
       return;
 
     case WebAssembly::DELEGATE:
     case WebAssembly::DELEGATE_S:
-      if (ControlFlowStack.empty() || EHPadStack.empty() ||
-          EHInstStack.empty()) {
+      if (ControlFlowStack.empty() || TryStack.empty() || EHInstStack.empty()) {
         printAnnotation(OS, "try-delegate mismatch!");
       } else {
         // 'delegate' is
         // 1. A marker for the end of block label
         // 2. A destination for throwing instructions
         // 3. An instruction that itself rethrows to another 'catch'
-        assert(ControlFlowStack.back().first == EHPadStack.back());
+        assert(ControlFlowStack.back().first == TryStack.back());
         std::string Label = "label/catch" +
                             utostr(ControlFlowStack.pop_back_val().first) +
                             ": ";
-        EHPadStack.pop_back();
+        TryStack.pop_back();
         EHInstStack.pop_back();
         uint64_t Depth = MI->getOperand(0).getImm();
         if (Depth >= ControlFlowStack.size()) {

llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyInstPrinter.h

@@ -26,7 +26,7 @@ class MCSubtargetInfo;
 class WebAssemblyInstPrinter final : public MCInstPrinter {
   uint64_t ControlFlowCounter = 0;
   SmallVector<std::pair<uint64_t, bool>, 4> ControlFlowStack;
-  SmallVector<uint64_t, 4> EHPadStack;
+  SmallVector<uint64_t, 4> TryStack;
 
   enum EHInstKind { TRY, CATCH, CATCH_ALL };
   SmallVector<EHInstKind, 4> EHInstStack;

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

@@ -86,6 +86,9 @@ class WebAssemblyCFGStackify final : public MachineFunctionPass {
                           const MachineBasicBlock *MBB);
   unsigned getDelegateDepth(const SmallVectorImpl<EndMarkerInfo> &Stack,
                             const MachineBasicBlock *MBB);
+  unsigned
+  getRethrowDepth(const SmallVectorImpl<EndMarkerInfo> &Stack,
+                  const SmallVectorImpl<const MachineBasicBlock *> &EHPadStack);
   void rewriteDepthImmediates(MachineFunction &MF);
   void fixEndsAtEndOfFunction(MachineFunction &MF);
   void cleanupFunctionData(MachineFunction &MF);
@@ -1551,9 +1554,48 @@ unsigned WebAssemblyCFGStackify::getDelegateDepth(
   return Depth;
 }
 
+unsigned WebAssemblyCFGStackify::getRethrowDepth(
+    const SmallVectorImpl<EndMarkerInfo> &Stack,
+    const SmallVectorImpl<const MachineBasicBlock *> &EHPadStack) {
+  unsigned Depth = 0;
+  // In our current implementation, rethrows always rethrow the exception caught
+  // by the innermost enclosing catch. This means while traversing Stack in the
+  // reverse direction, when we encounter END_TRY, we should check if the
+  // END_TRY corresponds to the current innermost EH pad. For example:
+  // try
+  //   ...
+  // catch         ;; (a)
+  //   try
+  //     rethrow 1 ;; (b)
+  //   catch       ;; (c)
+  //     rethrow 0 ;; (d)
+  //   end         ;; (e)
+  // end           ;; (f)
+  //
+  // When we are at 'rethrow' (d), while reversely traversing Stack the first
+  // 'end' we encounter is the 'end' (e), which corresponds to the 'catch' (c).
+  // And 'rethrow' (d) rethrows the exception caught by 'catch' (c), so we stop
+  // there and the depth should be 0. But when we are at 'rethrow' (b), it
+  // rethrows the exception caught by 'catch' (a), so when traversing Stack
+  // reversely, we should skip the 'end' (e) and choose 'end' (f), which
+  // corresponds to 'catch' (a).
+  for (auto X : reverse(Stack)) {
+    const MachineInstr *End = X.second;
+    if (End->getOpcode() == WebAssembly::END_TRY) {
+      auto *EHPad = TryToEHPad[EndToBegin[End]];
+      if (EHPadStack.back() == EHPad)
+        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;
+  SmallVector<const MachineBasicBlock *, 8> EHPadStack;
   for (auto &MBB : reverse(MF)) {
     for (auto I = MBB.rbegin(), E = MBB.rend(); I != E; ++I) {
       MachineInstr &MI = *I;
@@ -1575,16 +1617,28 @@ void WebAssemblyCFGStackify::rewriteDepthImmediates(MachineFunction &MF) {
         Stack.push_back(std::make_pair(&MBB, &MI));
         break;
 
-      case WebAssembly::END_TRY:
+      case WebAssembly::END_TRY: {
         // We handle DELEGATE in the default level, because DELEGATE has
         // immediate operands to rewrite.
         Stack.push_back(std::make_pair(&MBB, &MI));
+        auto *EHPad = TryToEHPad[EndToBegin[&MI]];
+        EHPadStack.push_back(EHPad);
         break;
+      }
 
       case WebAssembly::END_LOOP:
         Stack.push_back(std::make_pair(EndToBegin[&MI]->getParent(), &MI));
         break;
 
+      case WebAssembly::CATCH:
+      case WebAssembly::CATCH_ALL:
+        EHPadStack.pop_back();
+        break;
+
+      case WebAssembly::RETHROW:
+        MI.getOperand(0).setImm(getRethrowDepth(Stack, EHPadStack));
+        break;
+
       default:
         if (MI.isTerminator()) {
           // Rewrite MBB operands to be depth immediates.

llvm/lib/Target/WebAssembly/WebAssemblyInstrControl.td

@@ -133,8 +133,8 @@ defm THROW : I<(outs), (ins event_op:$tag, variable_ops),
                "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

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

@@ -119,7 +119,7 @@ try.cont:                                         ; preds = %catch, %catch2, %en
 ; CHECK:               rethrow   0                     # down to catch[[C0:[0-9]+]]
 ; CHECK:             end_try
 ; CHECK:           end_block                           # label[[L2]]:
-; CHECK:           rethrow   0                         # down to catch[[C0]]
+; CHECK:           rethrow   1                         # down to catch[[C0]]
 ; CHECK:         catch_all                             # catch[[C0]]:
 ; CHECK:           call      __cxa_end_catch
 ; CHECK:           rethrow   0                         # to caller
@@ -128,7 +128,7 @@ try.cont:                                         ; preds = %catch, %catch2, %en
 ; 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

llvm/test/CodeGen/WebAssembly/exception.mir

@@ -12,6 +12,9 @@
   define void @unreachable_ehpad_test() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
     ret void
   }
+  define void @rethrow_arg_test() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+    ret void
+  }
 ...
 
 ---
@@ -24,6 +27,7 @@ liveins:
 body: |
   bb.0:
     ; TRY should be before EH_LABEL wrappers of throwing calls
+    ; CHECK:      bb.0
     ; CHECK:      TRY
     ; CHECK-NEXT: EH_LABEL
     ; CHECK-NEXT: CALL @foo
@@ -38,15 +42,17 @@ body: |
   ; predecessors: %bb.0
     successors: %bb.2
     ; CATCH_ALL should be after an EH_LABEL at the beginning of an EH pad
+    ; CHECK:      bb.1
     ; CHECK:      EH_LABEL
     ; CHECK-NEXT: CATCH_ALL
     EH_LABEL <mcsymbol .Ltmp2>
-    CATCHRET %bb.2, %bb.0, implicit-def dead $arguments
+    CATCHRET %bb.2, %bb.1, implicit-def dead $arguments
 
   bb.2:
   ; predecessors: %bb.0, %bb.1
     RETURN implicit-def dead $arguments
 ...
+
 ---
 # Unreachable EH pads should be removed by LateEHPrepare.
 # CHECK-LABEL: name: unreachable_ehpad_test
@@ -64,10 +70,63 @@ body: |
   bb.1 (landing-pad):
     successors: %bb.2
     EH_LABEL <mcsymbol .Ltmp2>
-    CATCHRET %bb.2, %bb.0, implicit-def dead $arguments
+    CATCHRET %bb.2, %bb.1, implicit-def dead $arguments
 
   ; CHECK: bb.2
   bb.2:
   ; predecessors: %bb.0, %bb.1
     RETURN implicit-def dead $arguments
 ...
+
+---
+# CHECK-LABEL: name: rethrow_arg_test
+name: rethrow_arg_test
+liveins:
+  - { reg: '$arguments' }
+body: |
+  bb.0:
+    successors: %bb.1, %bb.4
+    ; CHECK: bb.0
+    ; CHECK: TRY
+    EH_LABEL <mcsymbol .Ltmp0>
+    CALL @foo, implicit-def dead $arguments, implicit $sp32, implicit $sp64
+    EH_LABEL <mcsymbol .Ltmp1>
+    BR %bb.4, implicit-def dead $arguments
+
+  bb.1 (landing-pad):
+  ; predecessors: %bb.0
+    successors: %bb.2
+    ; CHECK: bb.1
+    ; CHECK: CATCH
+    ; CHECK: TRY
+    ; This RETHROW rethrows the exception caught by this BB's CATCH, but after
+    ; CFGStackify a TRY is placed between the CATCH and this RETHROW, so after
+    ; CFGStackify its immediate argument should become not 0, but 1.
+    ; CHECK: RETHROW 1
+    EH_LABEL <mcsymbol .Ltmp2>
+    %0:i32 = CATCH &__cpp_exception, implicit-def dead $arguments
+    RETHROW 0, implicit-def dead $arguments
+
+  bb.2 (landing-pad):
+  ; predecessors: %bb.1
+    successors: %bb.3
+    ; CHECK: bb.2
+    ; CHECK: CATCH
+    ; CHECK: RETHROW 0
+    EH_LABEL <mcsymbol .Ltmp5>
+    %1:i32 = CATCH &__cpp_exception, implicit-def dead $arguments
+    RETHROW 0, implicit-def dead $arguments
+    CATCHRET %bb.3, %bb.2, implicit-def dead $arguments
+
+  bb.3:
+  ; predecessors: %bb.2
+    successors: %bb.4
+    CATCHRET %bb.4, %bb.1, implicit-def dead $arguments
+
+  bb.4:
+  ; predecessors: %bb.0, %bb.3
+    ; CHECK: bb.4
+    ; CHECK: END_TRY
+    ; CHECK: END_TRY
+    RETURN implicit-def dead $arguments
+