[WebAssembly] Fix delegate's argument computation

[aheejin]

Summary

I previously assumed delegate's immediate argument computation
followed a different rule than that of branches, but we agreed to make
it the same
(https://github.com/WebAssembly/exception-handling/issues/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.

[aheejin]

Remove a redundant line

[aheejin]

Fix the error message

[aheejin]

Use EndMarkerInfo

[tlively]

Nice simplification!

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

 namespace llvm {
 
 class MCSubtargetInfo;
 
 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;
 
 public:
   WebAssemblyInstPrinter(const MCAsmInfo &MAI, const MCInstrInfo &MII,
                          const MCRegisterInfo &MRI);
 

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

if (CommentStream) {
     unsigned Opc = MI->getOpcode();
     switch (Opc) {
     default:
       break;
 
     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;
 
     case WebAssembly::END_TRY:
     case WebAssembly::END_TRY_S:
       if (ControlFlowStack.empty() || EHInstStack.empty()) {
         printAnnotation(OS, "End marker mismatch!");
       } else {
if (CommentStream) {
     case WebAssembly::CATCH_ALL_S:
       // There can be multiple catch instructions for one try instruction, so
       // we print a label only for the first 'catch' label.
       if (EHInstStack.empty()) {
         printAnnotation(OS, "try-catch mismatch!");
       } 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) {
           EHInstStack.push_back(CATCH);
         } else {
           EHInstStack.push_back(CATCH_ALL);
         }
       }
       return;
 
     case WebAssembly::RETHROW:
     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()) {
         printAnnotation(OS, "to caller");
       } else {
         printAnnotation(OS, "down to catch" + utostr(EHPadStack.back()));
       }
       return;
 
     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);
       }
       return;
     }
 
     // Annotate any control flow label references.
 

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

class WebAssemblyCFGStackify final : public MachineFunctionPass {
   bool fixCallUnwindMismatches(MachineFunction &MF);
   bool fixCatchUnwindMismatches(MachineFunction &MF);
   void addTryDelegate(MachineInstr *RangeBegin, MachineInstr *RangeEnd,
                       MachineBasicBlock *DelegateDest);
   void recalculateScopeTops(MachineFunction &MF);
   void removeUnnecessaryInstrs(MachineFunction &MF);
 
   // Wrap-up
-  unsigned getDepth(const SmallVectorImpl<const MachineBasicBlock *> &Stack,
+  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);
 
   // 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 &MI : reverse(MBB)) {
       case WebAssembly::CATCH_ALL:
         updateScopeTops(EHPadToTry[&MBB]->getParent(), &MBB);
         break;
       }
     }
   }
 }
 
-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
 /// the function, because it's never reached. However, in WebAssembly, blocks
 /// that end at the function end need to have a return type signature that
 /// matches the function signature, even though it's unreachable. This function
 /// checks for such cases and fixes up the signatures.
 void WebAssemblyCFGStackify::fixEndsAtEndOfFunction(MachineFunction &MF) {
if (MCAI->getExceptionHandlingType() == ExceptionHandling::Wasm &&
       MF.getFunction().hasPersonalityFn()) {
     bool Changed = fixCallUnwindMismatches(MF);
     Changed |= fixCatchUnwindMismatches(MF);
     if (Changed)
       recalculateScopeTops(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() && "Branch 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<std::pair<const MachineBasicBlock *, const MachineInstr *>, 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:
         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()) {
               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;
       }
     }
   }
   assert(Stack.empty() && "Control flow should be balanced");
 }
 
 void WebAssemblyCFGStackify::cleanupFunctionData(MachineFunction &MF) {

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

 ; NOSORT:         call  bar
 ; NOSORT:       delegate    2            # label/catch{{[0-9]+}}: down to catch[[C0:[0-9]+]]
 ; --- try-delegate ends (call unwind mismatch)
 ; NOSORT:     catch
 ; NOSORT:       call  {{.*}} __cxa_begin_catch
 ; --- 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
 ; --- try-delegate ends (catch unwind mismatch)
 ; NOSORT: catch  {{.*}}                  # catch[[C0]]:
 ; NOSORT:   call  {{.*}} __cxa_begin_catch
 ; NOSORT:   call  __cxa_end_catch
 ; NOSORT: end_try