1 files changed, 65 insertions, 40 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
index 87e5f8e..745b97f 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
@@ -80,8 +80,12 @@ class WebAssemblyCFGStackify final : public MachineFunctionPass {
   void removeUnnecessaryInstrs(MachineFunction &MF);
 
   // Wrap-up
-  unsigned getDepth(const SmallVectorImpl<const MachineBasicBlock *> &Stack,
-                    const MachineBasicBlock *MBB);
+  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);
   void rewriteDepthImmediates(MachineFunction &MF);
   void fixEndsAtEndOfFunction(MachineFunction &MF);
   void cleanupFunctionData(MachineFunction &MF);
@@ -1399,21 +1403,6 @@ void WebAssemblyCFGStackify::recalculateScopeTops(MachineFunction &MF) {
   }
 }
 
-unsigned WebAssemblyCFGStackify::getDepth(
-    const SmallVectorImpl<const MachineBasicBlock *> &Stack,
-    const MachineBasicBlock *MBB) {
-  if (MBB == FakeCallerBB)
-    return Stack.size();
-  unsigned Depth = 0;
-  for (auto X : reverse(Stack)) {
-    if (X == MBB)
-      break;
-    ++Depth;
-  }
-  assert(Depth < Stack.size() && "Branch destination should be in scope");
-  return Depth;
-}
-
 /// In normal assembly languages, when the end of a function is unreachable,
 /// because the function ends in an infinite loop or a noreturn call or similar,
 /// it isn't necessary to worry about the function return type at the end of
@@ -1515,48 +1504,85 @@ void WebAssemblyCFGStackify::placeMarkers(MachineFunction &MF) {
   }
 }
 
+unsigned WebAssemblyCFGStackify::getBranchDepth(
+    const SmallVectorImpl<EndMarkerInfo> &Stack, const MachineBasicBlock *MBB) {
+  unsigned Depth = 0;
+  for (auto X : reverse(Stack)) {
+    if (X.first == MBB)
+      break;
+    ++Depth;
+  }
+  assert(Depth < Stack.size() && "Branch destination should be in scope");
+  return Depth;
+}
+
+unsigned WebAssemblyCFGStackify::getDelegateDepth(
+    const SmallVectorImpl<EndMarkerInfo> &Stack, const MachineBasicBlock *MBB) {
+  if (MBB == FakeCallerBB)
+    return Stack.size();
+  // Delegate's destination is either a catch or a another delegate BB. When the
+  // destination is another delegate, we can compute the argument in the same
+  // way as branches, because the target delegate BB only contains the single
+  // delegate instruction.
+  if (!MBB->isEHPad()) // Target is a delegate BB
+    return getBranchDepth(Stack, MBB);
+
+  // When the delegate's destination is a catch BB, we need to use its
+  // corresponding try's end_try BB because Stack contains each marker's end BB.
+  // Also we need to check if the end marker instruction matches, because a
+  // single BB can contain multiple end markers, like this:
+  // bb:
+  //   END_BLOCK
+  //   END_TRY
+  //   END_BLOCK
+  //   END_TRY
+  //   ...
+  //
+  // In case of branches getting the immediate that targets any of these is
+  // fine, but delegate has to exactly target the correct try.
+  unsigned Depth = 0;
+  const MachineInstr *EndTry = BeginToEnd[EHPadToTry[MBB]];
+  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;
+}
+
 void WebAssemblyCFGStackify::rewriteDepthImmediates(MachineFunction &MF) {
   // Now rewrite references to basic blocks to be depth immediates.
-  SmallVector<const MachineBasicBlock *, 8> Stack;
-  SmallVector<const MachineBasicBlock *, 8> DelegateStack;
+  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:
       case WebAssembly::TRY:
-        assert(ScopeTops[Stack.back()->getNumber()]->getNumber() <=
+        assert(ScopeTops[Stack.back().first->getNumber()]->getNumber() <=
                    MBB.getNumber() &&
                "Block/try marker should be balanced");
         Stack.pop_back();
-        DelegateStack.pop_back();
         break;
 
       case WebAssembly::LOOP:
-        assert(Stack.back() == &MBB && "Loop top should be balanced");
+        assert(Stack.back().first == &MBB && "Loop top should be balanced");
         Stack.pop_back();
-        DelegateStack.pop_back();
         break;
 
       case WebAssembly::END_BLOCK:
-        Stack.push_back(&MBB);
-        DelegateStack.push_back(&MBB);
+        Stack.push_back(std::make_pair(&MBB, &MI));
         break;
 
       case WebAssembly::END_TRY:
         // We handle DELEGATE in the default level, because DELEGATE has
-        // immediate operands to rewirte.
-        Stack.push_back(&MBB);
+        // immediate operands to rewrite.
+        Stack.push_back(std::make_pair(&MBB, &MI));
         break;
 
       case WebAssembly::END_LOOP:
-        Stack.push_back(EndToBegin[&MI]->getParent());
-        DelegateStack.push_back(EndToBegin[&MI]->getParent());
-        break;
-
-      case WebAssembly::CATCH:
-      case WebAssembly::CATCH_ALL:
-        DelegateStack.push_back(&MBB);
+        Stack.push_back(std::make_pair(EndToBegin[&MI]->getParent(), &MI));
         break;
 
       default:
@@ -1569,18 +1595,17 @@ void WebAssemblyCFGStackify::rewriteDepthImmediates(MachineFunction &MF) {
             if (MO.isMBB()) {
               if (MI.getOpcode() == WebAssembly::DELEGATE)
                 MO = MachineOperand::CreateImm(
-                    getDepth(DelegateStack, MO.getMBB()));
+                    getDelegateDepth(Stack, MO.getMBB()));
               else
-                MO = MachineOperand::CreateImm(getDepth(Stack, MO.getMBB()));
+                MO = MachineOperand::CreateImm(
+                    getBranchDepth(Stack, MO.getMBB()));
             }
             MI.addOperand(MF, MO);
           }
         }
 
-        if (MI.getOpcode() == WebAssembly::DELEGATE) {
-          Stack.push_back(&MBB);
-          DelegateStack.push_back(&MBB);
-        }
+        if (MI.getOpcode() == WebAssembly::DELEGATE)
+          Stack.push_back(std::make_pair(&MBB, &MI));
         break;
       }
     }