[WebAssembly] Fix delegate's argument computation

I previously assumed `delegate`'s immediate argument computation
followed a different rule than that of branches, but we agreed to make
it the same
(WebAssembly/exception-handling#146). This
removes the need for a separate `DelegateStack` in both CFGStackify and
InstPrinter.

When computing the immediate argument, we use a different function for
`delegate` computation because in MIR `DELEGATE`'s instruction's
destination is the destination catch BB or delegate BB, and when it is a
catch BB, we need an additional step of getting its corresponding `end`
marker.

Reviewed By: tlively, dschuff

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

Author: aheejin

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

@@ -93,42 +93,37 @@ void WebAssemblyInstPrinter::printInst(const MCInst *MI, uint64_t Address,
     case WebAssembly::LOOP:
     case WebAssembly::LOOP_S:
       printAnnotation(OS, "label" + utostr(ControlFlowCounter) + ':');
-      ControlFlowStack.push_back(std::make_pair(ControlFlowCounter, true));
-      DelegateStack.push_back(ControlFlowCounter++);
+      ControlFlowStack.push_back(std::make_pair(ControlFlowCounter++, true));
       return;
 
     case WebAssembly::BLOCK:
     case WebAssembly::BLOCK_S:
-      ControlFlowStack.push_back(std::make_pair(ControlFlowCounter, false));
-      DelegateStack.push_back(ControlFlowCounter++);
+      ControlFlowStack.push_back(std::make_pair(ControlFlowCounter++, false));
       return;
 
     case WebAssembly::TRY:
     case WebAssembly::TRY_S:
       ControlFlowStack.push_back(std::make_pair(ControlFlowCounter, false));
-      EHPadStack.push_back(ControlFlowCounter);
-      DelegateStack.push_back(ControlFlowCounter++);
+      EHPadStack.push_back(ControlFlowCounter++);
       EHInstStack.push_back(TRY);
       return;
 
     case WebAssembly::END_LOOP:
     case WebAssembly::END_LOOP_S:
-      if (ControlFlowStack.empty() || DelegateStack.empty()) {
+      if (ControlFlowStack.empty()) {
         printAnnotation(OS, "End marker mismatch!");
       } else {
         ControlFlowStack.pop_back();
-        DelegateStack.pop_back();
       }
       return;
 
     case WebAssembly::END_BLOCK:
     case WebAssembly::END_BLOCK_S:
-      if (ControlFlowStack.empty() || DelegateStack.empty()) {
+      if (ControlFlowStack.empty()) {
         printAnnotation(OS, "End marker mismatch!");
       } else {
         printAnnotation(
             OS, "label" + utostr(ControlFlowStack.pop_back_val().first) + ':');
-        DelegateStack.pop_back();
       }
       return;
 
@@ -154,12 +149,11 @@ 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() || DelegateStack.empty()) {
+        if (EHPadStack.empty()) {
           printAnnotation(OS, "try-catch mismatch!");
         } else {
           printAnnotation(OS,
                           "catch" + utostr(EHPadStack.pop_back_val()) + ':');
-          DelegateStack.pop_back();
         }
         EHInstStack.pop_back();
         if (Opc == WebAssembly::CATCH || Opc == WebAssembly::CATCH_S) {
@@ -184,26 +178,28 @@ void WebAssemblyInstPrinter::printInst(const MCInst *MI, uint64_t Address,
     case WebAssembly::DELEGATE:
     case WebAssembly::DELEGATE_S:
       if (ControlFlowStack.empty() || EHPadStack.empty() ||
-          DelegateStack.empty() || EHInstStack.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() &&
-               EHPadStack.back() == DelegateStack.back());
+        assert(ControlFlowStack.back().first == EHPadStack.back());
         std::string Label = "label/catch" +
                             utostr(ControlFlowStack.pop_back_val().first) +
                             ": ";
         EHPadStack.pop_back();
-        DelegateStack.pop_back();
         EHInstStack.pop_back();
         uint64_t Depth = MI->getOperand(0).getImm();
-        if (Depth >= DelegateStack.size()) {
+        if (Depth >= ControlFlowStack.size()) {
           Label += "to caller";
         } else {
-          Label += "down to catch" + utostr(DelegateStack.rbegin()[Depth]);
+          const auto &Pair = ControlFlowStack.rbegin()[Depth];
+          if (Pair.second)
+            printAnnotation(OS, "delegate cannot target a loop");
+          else
+            Label += "down to catch" + utostr(Pair.first);
         }
         printAnnotation(OS, Label);
       }

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

@@ -27,11 +27,6 @@ class WebAssemblyInstPrinter final : public MCInstPrinter {
   uint64_t ControlFlowCounter = 0;
   SmallVector<std::pair<uint64_t, bool>, 4> ControlFlowStack;
   SmallVector<uint64_t, 4> EHPadStack;
-  // 'delegate' can target any block-like structure, but in case the target is
-  // 'try', it rethrows to the corresponding 'catch'. Because it can target all
-  // blocks but with a slightly different semantics with branches, we need a
-  // separate stack for 'delegate'.
-  SmallVector<uint64_t, 4> DelegateStack;
 
   enum EHInstKind { TRY, CATCH, CATCH_ALL };
   SmallVector<EHInstKind, 4> EHInstStack;

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;
       }
     }

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

@@ -658,7 +658,7 @@ try.cont:                                         ; preds = %catch.start0
 ; --- try-delegate starts (call unwind mismatch)
 ; NOSORT:       try
 ; NOSORT:         call  __cxa_end_catch
-; NOSORT:       delegate    2            # label/catch{{[0-9]+}}: to caller
+; NOSORT:       delegate    3            # label/catch{{[0-9]+}}: to caller
 ; --- try-delegate ends (call unwind mismatch)
 ; NOSORT:     end_try
 ; NOSORT:   delegate    1                # label/catch{{[0-9]+}}: to caller