1 files changed, 55 insertions, 1 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
index 745b97f..9f5290a 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
+++ b/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.