[WebAssembly] Fix unwind destination mismatches in CFG stackify

[aheejin]

Summary

Linearing the control flow by placing try/end_try markers can create
mismatches in unwind destinations. This patch resolves these mismatches
by wrapping those instructions with an incorrect unwind destination with
a nested try/catch/end_try and branching to the right destination
within the new catch block.

[aheejin]

Continuing the discussion started in the email:

I think it's better not to split this pass as a separate pass. The reasons are:

1. Now CFGStackify pass maintains several internal maps (ScopeTops, BeginToEnd, EndToBegin, TryToEHPad, EHPadToTry, and BeginToBottom) to keep track of start/end markers, and entries are added to these maps and queried both in CFGStackify and namely, FixUnwindMismatch pass. So these maps have to be maintained in both passes, and also, we have to either carry that information using auxiliary data structure or we have to recalculate these maps at the end of FixUnwindMatch, which is more work.

2. Now CFGStackify also has several helper functions (registerScope(), registerTryScope(), and unregisterScope()) that help add and delete entries to various maps I mentioned in 1. We need these utility function in both passes.

3. rewriteDepthImmediate() and fixEndsAtEndOfFunction() functions run at the end of CFGStackify. But because we change CFG in FixUnwindMatch pass as well, they should rather run at the end of FixUnwindMatch pass now. This looks a bit weird bcause FixUnwindMatch pass should only run for functions that have a personality function (or an EH pad). And also intuitively rewriting depth immediates is not a part of FixUnwindMatch but CFGStackify.

Also, if we can bear with CFGStackify file getting a bit longer, I think the two parts are kind of well divided and factorized into functions.

[aheejin]

• Add unregisterScope() deleted in rL348647

[aheejin]

Reimplementation based on the new proposal

[aheejin]

Add comments in tests

[dschuff]

lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

75

This comment is pretty vague. I guess this is one block for each function, and is now shared for creating a correct signature for fallthrough returns, and to use as a target for rethrows that need to unwind to the caller, but are trapped inside an outer try/catch?

694

s/Linearing/Linearizing

730

s/boy/body

764

remove the period after "block"

834

It seems unfortunate that we have to care about whether explicit locals is turned on. Is that just because this pass runs after createWebAssemblyExplicitLocals? Does the order matter? e.g. would it be simpler to run it after CfgStackify? Or, could we just run it again if this pass changes anything?

869

continue here will just skip to the next instruction, but the comment says it will skip the rest of the BB (which makes sense), should this be break instead? Or do we need to keep going to find the try/catch?

943

s/don't/don't have

1012

these could be else if?

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

42

Any weird corner cases it would be worthwhile to test for? Mixing both kinds of mismatches? A case with uwind-to-caller and block signature needed? Would multiple nesting of mismatches affect what happens?

[aheejin]

Addressed simple comments first.
Things I haven't addressed yet:

• Add more test cases
• Switch ExplicitLocal pass order?

lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

75

Yes. Updated the comments.

869

Right. It can't be break because we still need to process this part and update EHPadStack.

if (MI.getOpcode() == WebAssembly::TRY)
  EHPadStack.pop_back();
else if (MI.getOpcode() == WebAssembly::CATCH)
  EHPadStack.push_back(MI.getParent());

1012

Changed to switch-case.

[aheejin]

lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

834

Actually I don't think the order matters. Thanks for the idea. Created D59652.

[aheejin]

Now after D59652 we don't need the separate code (and separate test commands) for when -wasm-disable-explicit-locals is set and it is not. Will fix it soon.

[aheejin]

• Deleted all logics related to -wasm-disable-explicit-locals thanks to D59652
• Fix several bugs
• Add more tests

[aheejin]

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

42

Add some more tests.

• Mixing both kinds of mismatches: Done (test8)
• A case with uwind-to-caller and block signature needed: it's not possible because of newly added rethrow. Added a test that shows why it's not possible. (test7)
• Multiple nesting of mismatches: test8 kind of covers this too.

Note that -wasm-disable-ehpad-sort is used to generate maximum # of mismatches for those mismatch tests.

[aheejin]

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

42

Oh and please also note that I deleted this file and moved all tests to cfg-stackify-eh.ll.

lib/Target/WebAssembly/WebAssemblyCFGSort.cpp

 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "wasm-cfg-sort"
 
+// Options to disable EH pad first sorting. Only for testing unwind destination
+// mismatches in CFGStackify.
+static cl::opt<bool> WasmDisableEHPadSort(
+    "wasm-disable-ehpad-sort", cl::ReallyHidden,
+    cl::desc(
+        "WebAssembly: Disable EH pad-first sort order. Testing purpose only."),
+    cl::init(false));
+
 namespace {
 
 // Wrapper for loops and exceptions
 class Region {
 public:
   virtual ~Region() = default;
   virtual MachineBasicBlock *getHeader() const = 0;
   virtual bool contains(const MachineBasicBlock *MBB) const = 0;

 //
 // If we give a higher priority to an EH pad whenever it is ready in this
 // example, when both bb1 and bb2 are ready, we would pick up bb2 first.
 
 /// Sort blocks by their number.
 struct CompareBlockNumbers {
   bool operator()(const MachineBasicBlock *A,
                   const MachineBasicBlock *B) const {
+    if (!WasmDisableEHPadSort) {
       if (A->isEHPad() && !B->isEHPad())
         return false;
       if (!A->isEHPad() && B->isEHPad())
         return true;
+    }
 
     return A->getNumber() > B->getNumber();
   }
 };
 /// Sort blocks by their number in the opposite order..
 struct CompareBlockNumbersBackwards {
   bool operator()(const MachineBasicBlock *A,
                   const MachineBasicBlock *B) const {
+    if (!WasmDisableEHPadSort) {
       // We give a higher priority to an EH pad
       if (A->isEHPad() && !B->isEHPad())
         return false;
       if (!A->isEHPad() && B->isEHPad())
         return true;
+    }
 
     return A->getNumber() < B->getNumber();
   }
 };
 /// Bookkeeping for a region to help ensure that we don't mix blocks not
 /// dominated by the its header among its blocks.
 struct Entry {
   const Region *TheRegion;

lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

 /// This is sufficient to convert arbitrary CFGs into a form that works on
 /// WebAssembly, provided that all loops are single-entry.
 ///
 /// In case we use exceptions, this pass also fixes mismatches in unwind
 /// destinations created during transforming CFG into wasm structured format.
 ///
 //===----------------------------------------------------------------------===//
 
-#include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
 #include "WebAssembly.h"
 #include "WebAssemblyExceptionInfo.h"
 #include "WebAssemblyMachineFunctionInfo.h"
 #include "WebAssemblySubtarget.h"
 #include "WebAssemblyUtilities.h"
 #include "llvm/CodeGen/MachineDominators.h"
-#include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
-#include "llvm/CodeGen/MachineLoopInfo.h"
-#include "llvm/CodeGen/MachineRegisterInfo.h"
-#include "llvm/CodeGen/Passes.h"
-#include "llvm/CodeGen/WasmEHFuncInfo.h"
 #include "llvm/MC/MCAsmInfo.h"
-#include "llvm/Support/Debug.h"
-#include "llvm/Support/raw_ostream.h"
-#include <cstring>
 using namespace llvm;
 
 #define DEBUG_TYPE "wasm-cfg-stackify"
 
+extern cl::opt<bool> WasmDisableExplicitLocals;
+
 namespace {
 class WebAssemblyCFGStackify final : public MachineFunctionPass {
   StringRef getPassName() const override { return "WebAssembly CFG Stackify"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<MachineDominatorTree>();
     AU.addRequired<MachineLoopInfo>();
     AU.addRequired<WebAssemblyExceptionInfo>();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
   // For each block whose label represents the end of a scope, record the block
   // which holds the beginning of the scope. This will allow us to quickly skip
   // over scoped regions when walking blocks.
   SmallVector<MachineBasicBlock *, 8> ScopeTops;
 
+  // Placing markers.
   void placeMarkers(MachineFunction &MF);
   void placeBlockMarker(MachineBasicBlock &MBB);
   void placeLoopMarker(MachineBasicBlock &MBB);
   void placeTryMarker(MachineBasicBlock &MBB);
   void removeUnnecessaryInstrs(MachineFunction &MF);
+  bool fixUnwindMismatches(MachineFunction &MF);
   void rewriteDepthImmediates(MachineFunction &MF);
   void fixEndsAtEndOfFunction(MachineFunction &MF);
 
   // For each BLOCK|LOOP|TRY, the corresponding END_(BLOCK|LOOP|TRY).
   DenseMap<const MachineInstr *, MachineInstr *> BeginToEnd;
   // For each END_(BLOCK|LOOP|TRY), the corresponding BLOCK|LOOP|TRY.
   DenseMap<const MachineInstr *, MachineInstr *> EndToBegin;
   // <TRY marker, EH pad> map
   DenseMap<const MachineInstr *, MachineBasicBlock *> TryToEHPad;
   // <EH pad, TRY marker> map
   DenseMap<const MachineBasicBlock *, MachineInstr *> EHPadToTry;
 
+  // Newly appended BB at the end of the function for special purpose
+  MachineBasicBlock *AppendixBB = nullptr;
+  MachineBasicBlock *getAppendixBlock(MachineFunction &MF) {
+    if (!AppendixBB) {
+      AppendixBB = MF.CreateMachineBasicBlock();
+      // Give it a fake predecessor so that AsmPrinter prints its label.
+      AppendixBB->addSuccessor(AppendixBB);
+      MF.push_back(AppendixBB);
+    }
+    return AppendixBB;
+  }
+
   // Helper functions to register / unregister scope information created by
   // marker instructions.
   void registerScope(MachineInstr *Begin, MachineInstr *End);
   void registerTryScope(MachineInstr *Begin, MachineInstr *End,
                         MachineBasicBlock *EHPad);
   void unregisterScope(MachineInstr *Begin);
 
 public:
   static char ID; // Pass identification, replacement for typeid
   WebAssemblyCFGStackify() : MachineFunctionPass(ID) {}
   ~WebAssemblyCFGStackify() override { releaseMemory(); }
   void releaseMemory() override;
 };
 } // end anonymous namespace
 
 char WebAssemblyCFGStackify::ID = 0;
 INITIALIZE_PASS(WebAssemblyCFGStackify, DEBUG_TYPE,
-                "Insert BLOCK and LOOP markers for WebAssembly scopes", false,
+                "Insert BLOCK/LOOP/TRY markers for WebAssembly scopes", false,
                 false)
 
 FunctionPass *llvm::createWebAssemblyCFGStackify() {
   return new WebAssemblyCFGStackify();
 }
 
 /// Test whether Pred has any terminators explicitly branching to MBB, as
 /// opposed to falling through. Note that it's possible (eg. in unoptimized
void WebAssemblyCFGStackify::unregisterScope(MachineInstr *Begin) {
   if (EHPad) {
     assert(EHPadToTry.count(EHPad));
     TryToEHPad.erase(Begin);
     EHPadToTry.erase(EHPad);
   }
 }
 
 /// Insert a BLOCK marker for branches to MBB (if needed).
+// TODO Consider a more generalized way of handling block (and also loop and
+// try) signatures when we implement the multi-value proposal later.
 void WebAssemblyCFGStackify::placeBlockMarker(MachineBasicBlock &MBB) {
   // This should have been handled in placeTryMarker.
   if (MBB.isEHPad())
     return;
 
   MachineFunction &MF = *MBB.getParent();
   auto &MDT = getAnalysis<MachineDominatorTree>();
   const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
void WebAssemblyCFGStackify::placeLoopMarker(MachineBasicBlock &MBB) {
   if (!Loop || Loop->getHeader() != &MBB)
     return;
 
   // The operand of a LOOP is the first block after the loop. If the loop is the
   // bottom of the function, insert a dummy block at the end.
   MachineBasicBlock *Bottom = WebAssembly::getBottom(Loop);
   auto Iter = std::next(Bottom->getIterator());
   if (Iter == MF.end()) {
-    MachineBasicBlock *Label = MF.CreateMachineBasicBlock();
-    // Give it a fake predecessor so that AsmPrinter prints its label.
-    Label->addSuccessor(Label);
-    MF.push_back(Label);
+    getAppendixBlock(MF);
     Iter = std::next(Bottom->getIterator());
   }
   MachineBasicBlock *AfterLoop = &*Iter;
 
   // Decide where in Header to put the LOOP.
   SmallPtrSet<const MachineInstr *, 4> BeforeSet;
   SmallPtrSet<const MachineInstr *, 4> AfterSet;
   for (const auto &MI : MBB) {
void WebAssemblyCFGStackify::placeTryMarker(MachineBasicBlock &MBB) {
   // If this try is at the bottom of the function, insert a dummy block at the
   // end.
   WebAssemblyException *WE = WEI.getExceptionFor(&MBB);
   assert(WE);
   MachineBasicBlock *Bottom = WebAssembly::getBottom(WE);
 
   auto Iter = std::next(Bottom->getIterator());
   if (Iter == MF.end()) {
-    MachineBasicBlock *Label = MF.CreateMachineBasicBlock();
-    // Give it a fake predecessor so that AsmPrinter prints its label.
-    Label->addSuccessor(Label);
-    MF.push_back(Label);
+    getAppendixBlock(MF);
     Iter = std::next(Bottom->getIterator());
   }
   MachineBasicBlock *Cont = &*Iter;
 
   assert(Cont != &MF.front());
   MachineBasicBlock *LayoutPred = Cont->getPrevNode();
 
   // If the nearest common dominator is inside a more deeply nested context,
void WebAssemblyCFGStackify::removeUnnecessaryInstrs(MachineFunction &MF) {
   }
   for (auto *MI : ToDelete) {
     if (MI->getOpcode() == WebAssembly::BLOCK)
       unregisterScope(MI);
     MI->eraseFromParent();
   }
 }
 
+bool WebAssemblyCFGStackify::fixUnwindMismatches(MachineFunction &MF) {
+  const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
+  auto &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
+  MachineRegisterInfo &MRI = MF.getRegInfo();
+
+  // Linearing the control flow by placing TRY / END_TRY markers can create
+  // mismatches in unwind destinations. There are two kinds of mismatches we
+  // try to solve here.
+
+  // 1. When an instruction may throw, but the EH pad it will unwind to can be
+  //    different from the original CFG.
+  //
+  // Example: we have the following CFG:
+  // bb0:
+  //   call @foo (if it throws, unwind to bb2)
+  // bb1:
+  //   call @bar (if it throws, unwind to bb3)
+  // bb2 (ehpad):
+  //   catch
+  //   ...
+  // bb3 (ehpad)
+  //   catch
+  //   handler body
+  //
+  // And the CFG is sorted in this order. Then after placing TRY markers, it
+  // will look like: (BB markers are omitted)
+  // try $label1
+  //   try
+  //     call @foo
+  //     call @bar   (if it throws, unwind to bb3)
+  //   catch         <- ehpad (bb2)
+  //     ...
+  //   end_try
+  // catch           <- ehpad (bb3)
+  //   handler body
+  // end_try
+  //
+  // Now if bar() throws, it is going to end up ip in bb2, not bb3, where it
+  // is supposed to end up. We solve this problem by
+  // a. Split the target unwind EH pad (here bb3) so that the handler body is
+  //    right after 'end_try', which means we extract the handler body out of
+  //    the catch block. We do this because this handler boy should be somewhere
+  //    branch-eable from the inner scope.
+  // b. Wrap the call that has an incorrect unwind destination ('call @bar'
+  //    here) with a nested try/catch/end_try scope, and within the new catch
+  //    block, branches to the handler body.
+  // c. Place a branch after the newly inserted nested end_try so it can bypass
+  //    the handler body, which is now outside of a catch block.
+  //
+  // The result will like as follows. (new: a) means this instruction is newly
+  // created in the process of doing 'a' above.
+  //
+  // block $label0            (new: placeBlockMarker)
+  //   try $label1
+  //     try
+  //       call @foo
+  //       try                     (new: b)
+  //         call @bar
+  //       catch                   (new: b)
+  //         local.set n / drop    (new: b)
+  //         br $label1            (new: b)
+  //       end_try                 (new: b)
+  //     catch                     <- ehpad (bb2)
+  //     end_try
+  //     br $label0                (new: c)
+  //   catch                       <- ehpad (bb3)
+  //   end_try                     (hoisted: a)
+  //   handler body
+  // end_block                     (new: placeBlockMarker)
+  //
+  // Note that the new wrapping block/end_block will be generated later in
+  // placeBlockMarker.
+  //
+  // TODO Currently local.set and local.gets are generated to move except_ref
+  // value created by catches. That's because we don't support yielding values
+  // from a block. in LLVM machine IR, even though it is supported by wasm.
+  // Delete unnecessary local.get/local.sets once yielding values from a block
+  // is supported. The full EH spec requires multi-value support to do this, but
+  // for C++ we don't yet need it because we only throw a single i32.
+  //
+  // ---
+  // 2. The same as 1, but in this case an instruction unwinds to a caller
+  //    function and not another EH pad.
+  //
+  // Example: we have the following CFG:
+  // bb0:
+  //   call @foo (if it throws, unwind to bb2)
+  // bb1:
+  //   call @bar (if it throws, unwind to caller)
+  // bb2 (ehpad):
+  //   catch
+  //   ...
+  //
+  // And the CFG is sorted in this order. Then after placing TRY markers, it
+  // will look like:
+  // try
+  //   call @foo
+  //   call @bar   (if it throws, unwind to caller)
+  // catch         <- ehpad (bb2)
+  //   ...
+  // end_try
+  //
+  // Now if bar() throws, it is going to end up ip in bb2, when it is supposed
+  // throw up to the caller.
+  // We solve this problem by
+  // a. Create a new 'appendix' BB at the end of the function and put a single
+  //    'rethrow' instruction (+ local.get) in there.
+  // b. Wrap the call that has an incorrect unwind destination ('call @bar'
+  //    here) with a nested try/catch/end_try scope, and within the new catch
+  //    block, branches to the new appendix block.
+  //
+  // block $label0          (new: placeBlockMarker)
+  //   try
+  //     call @foo
+  //     try                (new: b)
+  //       call @bar
+  //     catch              (new: b)
+  //       local.set n      (new: b)
+  //       br $label0       (new: b)
+  //     end_try            (new: b)
+  //   catch                <- ehpad (bb2)
+  //     ...
+  //   end_try
+  // ...
+  // end_block              (new: placeBlockMarker)
+  // local.get n            (new: a)  <- appendix block
+  // rethrow                (new: a)
+  //
+  // In case there are multiple calls in a BB that may throw to the caller, they
+  // can be wrapped together in one nested try scope. (In 1, this couldn't
+  // happen, because may-throwing instruction there had an unwind destination,
+  // i.e., it was an invoke before, and there could be only one invoke within a
+  // BB.)
+
+  SmallVector<const MachineBasicBlock *, 8> EHPadStack;
+  // Range of intructions to be wrapped in a new nested try/catch
+  using TryRange = std::pair<MachineInstr *, MachineInstr *>;
+  // In original CFG, <unwind destionation BB, a vector of try ranges>
+  DenseMap<MachineBasicBlock *, SmallVector<TryRange, 4>> UnwindDestToTryRanges;
+  // In new CFG, <destination to branch to, a vector of try ranges>
+  DenseMap<MachineBasicBlock *, SmallVector<TryRange, 4>> BrDestToTryRanges;
+
+  // Information to be used when carrying except_ref value with branches. We
+  // currently do not yet support blocks returning values in the toolchain, and
+  // there can be multiple cases:
+  // 1. We have -wasm-disable-explicit-locals set.
+  // 2. We don't have -wasm-disable-explicit-locals set.
+  //   a. We need a 'local.set' after newly generated 'catch'.
+  //   b. We need a 'drop' after newly generated 'catch'.
+  // TODO Improve this after we have multi-value proposal support
+  struct ExnInfo {
+    // If we don't have -wasm-disable-explicit-locals, we need either a
+    // local.set or a drop. We will set one of them to true later.
+    bool NeedLocalSet = false;
+    bool NeedDrop = false;
+    // If -wasm-disable-explicit-locals is set, we just copy a register.
+    bool NeedCopy = WasmDisableExplicitLocals;
+    unsigned Local = 0; // Local to hold this except_ref value
+    unsigned Reg = 0;   // Register to hold this except_ref value
+  };
+
+  // <branch destination, ExnValInfo>
+  DenseMap<MachineBasicBlock *, ExnInfo> BrDestToExnInfo;
+
+  // Gather possibly throwing calls (i.e., previously invokes) whose current
+  // unwind destination is not the same as the original CFG.
+  for (auto &MBB : reverse(MF)) {
+    bool SeenThrowableInstInBB = false;
+    for (auto &MI : reverse(MBB)) {
+      if (MI.getOpcode() == WebAssembly::TRY)
+        EHPadStack.pop_back();
+      else if (MI.getOpcode() == WebAssembly::CATCH)
+        EHPadStack.push_back(MI.getParent());
+
+      // In this loop we only gather calls that have an EH pad to unwind. So
+      // there will be at most 1 such call (= invoke) in a BB, so after we've
+      // seen one, we can skip the rest of BB. Also if MBB has no EH pad
+      // successor or MI does not throw, this is not an invoke.
+      if (SeenThrowableInstInBB || !MBB.hasEHPadSuccessor() ||
+          !WebAssembly::mayThrow(MI))
+        continue;
+      SeenThrowableInstInBB = true;
+
+      // If the EH pad on the stack top is where this instruction should unwind
+      // next, we're good.
+      MachineBasicBlock *UnwindDest = nullptr;
+      for (auto *Succ : MBB.successors()) {
+        if (Succ->isEHPad()) {
+          UnwindDest = Succ;
+          break;
+        }
+      }
+      if (EHPadStack.back() == UnwindDest)
+        continue;
+
+      // If not, record the range.
+      UnwindDestToTryRanges[UnwindDest].push_back(TryRange(&MI, &MI));
+    }
+  }
+
+  assert(EHPadStack.empty());
+
+  // Gather possibly throwing calls that are supposed to unwind up to the caller
+  // if they throw, but currently unwind to an incorrect destination. Unlike the
+  // loop above, there can be multiple calls within a BB that unwind to the
+  // caller, which we should group together in a range.
+  bool NeedAppendixBlock = false;
+  for (auto &MBB : reverse(MF)) {
+    MachineInstr *RangeBegin = nullptr, *RangeEnd = nullptr; // inclusive
+    for (auto &MI : reverse(MBB)) {
+      if (MI.getOpcode() == WebAssembly::TRY)
+        EHPadStack.pop_back();
+      else if (MI.getOpcode() == WebAssembly::CATCH)
+        EHPadStack.push_back(MI.getParent());
+
+      // If MBB has an EH pad successor, this inst does not unwind to caller.
+      if (MBB.hasEHPadSuccessor())
+        continue;
+
+      // We wrap up the current range when we see a marker even if we haven't
+      // finished a BB.
+      if (RangeEnd && WebAssembly::isMarker(MI)) {
+        NeedAppendixBlock = true;
+        // Record the range. nullptr here means the unwind destination is the
+        // caller.
+        UnwindDestToTryRanges[nullptr].push_back(
+            TryRange(RangeBegin, RangeEnd));
+        RangeBegin = RangeEnd = nullptr; // Reset range pointers
+      }
+
+      // If EHPadStack is empty, that means it is correctly unwind to caller if
+      // it throws, so we're good. If MI does not throw, we're good too.
+      if (EHPadStack.empty() || !WebAssembly::mayThrow(MI))
+        continue;
+
+      // We found an instruction that unwinds to the caller but currently has an
+      // incorrect unwind destination. Create a new range or increment the
+      // currently existing range.
+      if (!RangeEnd)
+        RangeBegin = RangeEnd = &MI;
+      else
+        RangeBegin = &MI;
+    }
+
+    if (RangeEnd) {
+      NeedAppendixBlock = true;
+      // Record the range. nullptr here means the unwind destination is the
+      // caller.
+      UnwindDestToTryRanges[nullptr].push_back(TryRange(RangeBegin, RangeEnd));
+      RangeBegin = RangeEnd = nullptr; // Reset range pointers
+    }
+  }
+
+  assert(EHPadStack.empty());
+  // We don't any unwind destination mismatches to resolve.
+  if (UnwindDestToTryRanges.empty())
+    return false;
+
+  // If we found instructions that should unwind to the caller but currently
+  // have incorrect unwind destination, we create an appendix block at the end
+  // of the function with a local.get and a rethrow instruction.
+  if (NeedAppendixBlock) {
+    auto *AppendixBB = getAppendixBlock(MF);
+    ExnInfo Info;
+    Info.Reg = MRI.createVirtualRegister(&WebAssembly::EXCEPT_REFRegClass);
+    if (!WasmDisableExplicitLocals) {
+      Info.Local = MFI.getNumLocals();
+      Info.NeedLocalSet = true;
+      MFI.addLocal(MVT::ExceptRef);
+      BuildMI(AppendixBB, DebugLoc(),
+              TII.get(WebAssembly::LOCAL_GET_EXCEPT_REF), Info.Reg)
+          .addImm(Info.Local);
+      MFI.stackifyVReg(Info.Reg);
+    }
+    BuildMI(AppendixBB, DebugLoc(), TII.get(WebAssembly::RETHROW))
+        .addReg(Info.Reg);
+    // These instruction ranges should branch to this appendix BB.
+    for (auto Range : UnwindDestToTryRanges[nullptr])
+      BrDestToTryRanges[AppendixBB].push_back(Range);
+    BrDestToExnInfo[AppendixBB] = Info;
+  }
+
+  // We loop through unwind destination EH pads that are targeted from some
+  // inner scopes. Because these EH pads are destination of more than one scope
+  // now, we split them so that the handler body is after 'end_try'.
+  // - Before
+  // ehpad:
+  //   catch
+  //   local.set n / drop
+  //   handler body
+  // ...
+  // cont:
+  //   end_try
+  //
+  // - After
+  // ehpad:
+  //   catch
+  //   local.set n / drop
+  // brdest:               (new)
+  //   end_try             (hoisted from 'cont' BB)
+  //   handler body        (taken from 'ehpad')
+  // ...
+  // cont:
+  for (auto &P : UnwindDestToTryRanges) {
+    // This means the destination is the appendix BB, which was separately
+    // handled above.
+    if (!P.first)
+      continue;
+
+    MachineBasicBlock *EHPad = P.first;
+
+    // Find 'catch' and 'local.set' or 'drop' instruction that follows the
+    // 'catch'. If -wasm-disable-explicit-locals is not set, 'catch' should be
+    // always followed by either 'local.set' or a 'drop', because 'br_on_exn' is
+    // generated after 'catch' in LateEHPrepare and we don't support blocks
+    // taking values yet.
+    ExnInfo Info;
+    MachineInstr *Catch = nullptr, *LastUse = nullptr;
+    for (auto &MI : *EHPad) {
+      if (MI.getOpcode() == WebAssembly::CATCH) {
+        Catch = LastUse = &MI;
+        Info.Reg = Catch->getOperand(0).getReg();
+      }
+      if (MI.getOpcode() == WebAssembly::LOCAL_SET_EXCEPT_REF) {
+        assert(Catch && MI.getOperand(1).getReg() == Info.Reg);
+        LastUse = &MI;
+        Info.Local = MI.getOperand(1).getReg();
+        Info.NeedLocalSet = true;
+      }
+      if (MI.getOpcode() == WebAssembly::DROP_EXCEPT_REF) {
+        assert(Catch && MI.getOperand(0).getReg() == Info.Reg);
+        LastUse = &MI;
+        Info.NeedDrop = true;
+      }
+    }
+    assert(Catch && "EH pad does not have a catch");
+    assert((!Info.NeedLocalSet || !Info.NeedDrop) &&
+           "There shouldn't be both local.set and drop for except_ref");
+
+    auto SplitPos = std::next(LastUse->getIterator());
+
+    // Create a new BB that's gonna be the destination for branches from the
+    // inner mismatched scope.
+    MachineInstr *BeginTry = EHPadToTry[EHPad];
+    MachineInstr *EndTry = BeginToEnd[BeginTry];
+    MachineBasicBlock *Cont = EndTry->getParent();
+    auto *BrDest = MF.CreateMachineBasicBlock();
+    MF.insert(std::next(EHPad->getIterator()), BrDest);
+    // Hoist up the existing 'end_try'.
+    BrDest->insert(BrDest->end(), EndTry->removeFromParent());
+    // Take out the handler body from EH pad to the new branch destination BB.
+    BrDest->splice(BrDest->end(), EHPad, SplitPos, EHPad->end());
+    // Fix predecessor-successor relationship.
+    BrDest->transferSuccessors(EHPad);
+    EHPad->addSuccessor(BrDest);
+
+    // All try ranges that were supposed to unwind to this EH pad now have to
+    // branch to this new branch dest BB.
+    for (auto Range : UnwindDestToTryRanges[EHPad])
+      BrDestToTryRanges[BrDest].push_back(Range);
+    BrDestToExnInfo[BrDest] = Info;
+
+    // In case we fall through to the continuation BB after the catch block, we
+    // now have to add a branch to it.
+    // - Before
+    // try
+    //   ...
+    //   (falls through to 'cont')
+    // catch
+    //   handler body
+    // end
+    //               <-- cont
+    //
+    // - After
+    // try
+    //   ...
+    //   br %cont    (new)
+    // catch
+    // end
+    // handler body
+    //               <-- cont
+    MachineBasicBlock *EHPadLayoutPred = &*std::prev(EHPad->getIterator());
+    MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
+    SmallVector<MachineOperand, 4> Cond;
+    bool Analyzable = !TII.analyzeBranch(*EHPadLayoutPred, TBB, FBB, Cond);
+    if (Analyzable && !TBB && !FBB) {
+      DebugLoc DL = EHPadLayoutPred->empty()
+                        ? DebugLoc()
+                        : EHPadLayoutPred->rbegin()->getDebugLoc();
+      BuildMI(EHPadLayoutPred, DL, TII.get(WebAssembly::BR)).addMBB(Cont);
+    }
+  }
+
+  // For possibly throwing calls whose unwind destinations are currently
+  // incorrect because of CFG linearization, we wrap them with a nested
+  // try/catch/end_try, and within the new catch block, we branch to the correct
+  // handler.
+  // - Before
+  // mbb:
+  //   call @foo       <- Unwind destination mismatch!
+  // ehpad:
+  //   ...
+  //
+  // - After
+  // mbb:
+  //   try                (new)
+  //   call @foo
+  // nested-ehpad:        (new)
+  //   catch              (new)
+  //   local.set n / drop (new)
+  //   br %brdest         (new)
+  // nested-end:          (new)
+  //   end_try            (new)
+  // ehpad:
+  //   ...
+  for (auto &P : BrDestToTryRanges) {
+    MachineBasicBlock *BrDest = P.first;
+    auto &TryRanges = P.second;
+    ExnInfo Info = BrDestToExnInfo[BrDest];
+
+    for (auto Range : TryRanges) {
+      MachineInstr *RangeBegin = nullptr, *RangeEnd = nullptr;
+      std::tie(RangeBegin, RangeEnd) = Range;
+      auto *MBB = RangeBegin->getParent();
+
+      // Include possible EH_LABELs in the range
+      if (RangeBegin->getIterator() != MBB->begin() &&
+          std::prev(RangeBegin->getIterator())->isEHLabel())
+        RangeBegin = &*std::prev(RangeBegin->getIterator());
+      if (std::next(RangeEnd->getIterator()) != MBB->end() &&
+          std::next(RangeEnd->getIterator())->isEHLabel())
+        RangeEnd = &*std::next(RangeEnd->getIterator());
+
+      MachineBasicBlock *EHPad = nullptr;
+      for (auto *Succ : MBB->successors()) {
+        if (Succ->isEHPad()) {
+          EHPad = Succ;
+          break;
+        }
+      }
+
+      // Create the nested try instruction.
+      MachineInstr *NestedTry =
+          BuildMI(*MBB, *RangeBegin, RangeBegin->getDebugLoc(),
+                  TII.get(WebAssembly::TRY))
+              .addImm(int64_t(WebAssembly::ExprType::Void));
+
+      // Create the nested EH pad and fill instructions in.
+      MachineBasicBlock *NestedEHPad = MF.CreateMachineBasicBlock();
+      MF.insert(std::next(MBB->getIterator()), NestedEHPad);
+      NestedEHPad->setIsEHPad();
+      NestedEHPad->setIsEHScopeEntry();
+      unsigned ExnReg =
+          MRI.createVirtualRegister(&WebAssembly::EXCEPT_REFRegClass);
+      BuildMI(NestedEHPad, RangeEnd->getDebugLoc(), TII.get(WebAssembly::CATCH),
+              ExnReg);
+      // Depending on the cases, we need different instructions.
+      if (Info.NeedLocalSet) {
+        BuildMI(NestedEHPad, RangeEnd->getDebugLoc(),
+                TII.get(WebAssembly::LOCAL_SET_EXCEPT_REF))
+            .addImm(BrDestToExnInfo[BrDest].Local)
+            .addReg(ExnReg);
+        MFI.stackifyVReg(ExnReg);
+      } else if (Info.NeedDrop) {
+        BuildMI(NestedEHPad, RangeEnd->getDebugLoc(),
+                TII.get(WebAssembly::DROP_EXCEPT_REF))
+            .addReg(ExnReg);
+        MFI.stackifyVReg(ExnReg);
+      } else {
+        assert(Info.NeedCopy);
+        BuildMI(NestedEHPad, RangeEnd->getDebugLoc(),
+                TII.get(WebAssembly::COPY_EXCEPT_REF), Info.Reg)
+            .addReg(ExnReg);
+      };
+      BuildMI(NestedEHPad, RangeEnd->getDebugLoc(), TII.get(WebAssembly::BR))
+          .addMBB(BrDest);
+
+      // Create the nested continuation BB and end_try instruction.
+      MachineBasicBlock *NestedCont = MF.CreateMachineBasicBlock();
+      MF.insert(std::next(NestedEHPad->getIterator()), NestedCont);
+      MachineInstr *NestedEndTry =
+          BuildMI(*NestedCont, NestedCont->begin(), RangeEnd->getDebugLoc(),
+                  TII.get(WebAssembly::END_TRY));
+      // In case MBB has more instructions after the try range, move them to the
+      // new nested continuation BB.
+      NestedCont->splice(NestedCont->end(), MBB,
+                         std::next(RangeEnd->getIterator()), MBB->end());
+      registerTryScope(NestedTry, NestedEndTry, NestedEHPad);
+
+      // Fix predecessor-successor relationship.
+      NestedCont->transferSuccessors(MBB);
+      if (EHPad)
+        NestedCont->removeSuccessor(EHPad);
+      MBB->addSuccessor(NestedEHPad);
+      NestedEHPad->addSuccessor(BrDest);
+    }
+  }
+
+  // Renumber BBs and recalculate ScopeTop info because new BBs might have been
+  // created and inserted above.
+  MF.RenumberBlocks();
+  ScopeTops.clear();
+  ScopeTops.resize(MF.getNumBlockIDs());
+  for (auto &MBB : reverse(MF)) {
+    for (auto &MI : reverse(MBB)) {
+      if (ScopeTops[MBB.getNumber()])
+        break;
+      switch (MI.getOpcode()) {
+      case WebAssembly::END_BLOCK:
+      case WebAssembly::END_LOOP:
+      case WebAssembly::END_TRY:
+        ScopeTops[MBB.getNumber()] = EndToBegin[&MI]->getParent();
+        break;
+      case WebAssembly::CATCH:
+        ScopeTops[MBB.getNumber()] = EHPadToTry[&MBB]->getParent();
+        break;
+      }
+    }
+  }
+
+  // Recompute the dominator tree.
+  getAnalysis<MachineDominatorTree>().runOnMachineFunction(MF);
+  return true;
+}
+
 static unsigned
 getDepth(const SmallVectorImpl<const MachineBasicBlock *> &Stack,
          const MachineBasicBlock *MBB) {
   unsigned Depth = 0;
   for (auto X : reverse(Stack)) {
     if (X == MBB)
       break;
     ++Depth;
void WebAssemblyCFGStackify::placeMarkers(MachineFunction &MF) {
   // accommodate for the possible fake block we may insert at the end.
   ScopeTops.resize(MF.getNumBlockIDs() + 1);
   // Place the LOOP for MBB if MBB is the header of a loop.
   for (auto &MBB : MF)
     placeLoopMarker(MBB);
   // Place the TRY for MBB if MBB is the EH pad of an exception.
   const MCAsmInfo *MCAI = MF.getTarget().getMCAsmInfo();
   if (MCAI->getExceptionHandlingType() == ExceptionHandling::Wasm &&
-      MF.getFunction().hasPersonalityFn())
+      MF.getFunction().hasPersonalityFn()) {
     for (auto &MBB : MF)
       placeTryMarker(MBB);
+    // Fix mismatches in unwind destinations induced by linearizing the code.
+    fixUnwindMismatches(MF);
+  }
   // Place the BLOCK for MBB if MBB is branched to from above.
   for (auto &MBB : MF)
     placeBlockMarker(MBB);
 }
 
 void WebAssemblyCFGStackify::rewriteDepthImmediates(MachineFunction &MF) {
   // Now rewrite references to basic blocks to be depth immediates.
   SmallVector<const MachineBasicBlock *, 8> Stack;

 }
 
 void WebAssemblyCFGStackify::releaseMemory() {
   ScopeTops.clear();
   BeginToEnd.clear();
   EndToBegin.clear();
   TryToEHPad.clear();
   EHPadToTry.clear();
+  AppendixBB = nullptr;
 }
 
 bool WebAssemblyCFGStackify::runOnMachineFunction(MachineFunction &MF) {
   LLVM_DEBUG(dbgs() << "********** CFG Stackifying **********\n"
                        "********** Function: "
                     << MF.getName() << '\n');
   const MCAsmInfo *MCAI = MF.getTarget().getMCAsmInfo();
 
   releaseMemory();
 
   // Liveness is not tracked for VALUE_STACK physreg.
   MF.getRegInfo().invalidateLiveness();
 
   // Place the BLOCK/LOOP/TRY markers to indicate the beginnings of scopes.
   placeMarkers(MF);
 
+  // Remove unnecessary instructions possibly introduced by try/end_trys.
   if (MCAI->getExceptionHandlingType() == ExceptionHandling::Wasm &&
       MF.getFunction().hasPersonalityFn())
-    // Remove unnecessary instructions.
     removeUnnecessaryInstrs(MF);
 
   // Convert MBB operands in terminators to relative depth immediates.
   rewriteDepthImmediates(MF);
 
   // Fix up block/loop/try signatures at the end of the function to conform to
   // WebAssembly's rules.
   fixEndsAtEndOfFunction(MF);

lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp

 using namespace llvm;
 
 #define DEBUG_TYPE "wasm-explicit-locals"
 
 // A command-line option to disable this pass, and keep implicit locals
 // for the purpose of testing with lit/llc ONLY.
 // This produces output which is not valid WebAssembly, and is not supported
 // by assemblers/disassemblers and other MC based tools.
-static cl::opt<bool> WasmDisableExplicitLocals(
+cl::opt<bool> WasmDisableExplicitLocals(
     "wasm-disable-explicit-locals", cl::Hidden,
     cl::desc("WebAssembly: output implicit locals in"
              " instruction output for test purposes only."),
     cl::init(false));
 
 namespace {
 class WebAssemblyExplicitLocals final : public MachineFunctionPass {
   StringRef getPassName() const override {

lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.h

public:
   const std::vector<MVT> &getResults() const { return Results; }
 
   void clearParamsAndResults() {
     Params.clear();
     Results.clear();
   }
 
   void setNumLocals(size_t NumLocals) { Locals.resize(NumLocals, MVT::i32); }
+  size_t getNumLocals() const { return Locals.size(); }
   void setLocal(size_t i, MVT VT) { Locals[i] = VT; }
   void addLocal(MVT VT) { Locals.push_back(VT); }
   const std::vector<MVT> &getLocals() const { return Locals; }
 
   unsigned getVarargBufferVreg() const {
     assert(VarargVreg != -1U && "Vararg vreg hasn't been set");
     return VarargVreg;
   }

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

+; RUN: llc < %s -disable-wasm-fallthrough-return-opt -wasm-keep-registers -disable-block-placement -verify-machineinstrs -fast-isel=false -machine-sink-split-probability-threshold=0 -cgp-freq-ratio-to-skip-merge=1000 -exception-model=wasm -mattr=+exception-handling -wasm-disable-ehpad-sort | FileCheck %s --check-prefixes CHECK,LOCAL
+; RUN: llc < %s -disable-wasm-fallthrough-return-opt -wasm-keep-registers -disable-block-placement -verify-machineinstrs -fast-isel=false -machine-sink-split-probability-threshold=0 -cgp-freq-ratio-to-skip-merge=1000 -exception-model=wasm -mattr=+exception-handling -wasm-disable-ehpad-sort -wasm-disable-explicit-locals | FileCheck %s --check-prefixes CHECK,NOLOCAL
+
+target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+target triple = "wasm32-unknown-unknown"
+
+@_ZTIi = external constant i8*
+@_ZTId = external constant i8*
+
+; CHECK-LABEL: test0
+; CHECK:   block
+; CHECK:     try
+; CHECK:       try
+; CHECK:         call      foo
+; --- Nested try/catch/end_try starts here
+; CHECK:         try
+; CHECK:           call      bar
+; LOCAL:         catch     $push[[EXN:[0-9]+]]=          # catch2:
+; LOCAL:           drop  $pop[[EXN]]
+; NOLOCAL:       catch     $[[REG:[0-9]+]]=              # catch2:
+; NOLOCAL:         local.copy  $drop=, $[[REG]]
+; CHECK:           br        2                           # 2: down to label1
+; CHECK:         end_try                                 # label3:
+; --- Nested try/catch/end_try ends here
+; CHECK:         br        2                             # 2: down to label0
+; CHECK:       catch     {{.*}}                          # catch1:
+; CHECK:         br        2                             # 2: down to label0
+; CHECK:       end_try                                   # label2:
+; CHECK:     catch     {{.*}}                            # catch0:
+; LOCAL:       drop  {{.*}}
+; CHECK:     end_try                                     # label1:
+; CHECK:   end_block                                     # label0:
+; CHECK:   return
+
+define void @test0() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+bb0:
+  invoke void @foo()
+          to label %bb1 unwind label %catch.dispatch0
+
+bb1:
+  invoke void @bar()
+          to label %try.cont unwind label %catch.dispatch1
+
+catch.dispatch0:                                  ; preds = %bb0
+  %0 = catchswitch within none [label %catch.start0] unwind to caller
+
+catch.start0:                                     ; preds = %catch.dispatch0
+  %1 = catchpad within %0 [i8* null]
+  %2 = call i8* @llvm.wasm.get.exception(token %1)
+  %3 = call i32 @llvm.wasm.get.ehselector(token %1)
+  catchret from %1 to label %try.cont
+
+catch.dispatch1:                                  ; preds = %bb1
+  %4 = catchswitch within none [label %catch.start1] unwind to caller
+
+catch.start1:                                     ; preds = %catch.dispatch1
+  %5 = catchpad within %4 [i8* null]
+  %6 = call i8* @llvm.wasm.get.exception(token %5)
+  %7 = call i32 @llvm.wasm.get.ehselector(token %5)
+  catchret from %5 to label %try.cont
+
+try.cont:                                         ; preds = %bb0, %catch.start0
+  ret void
+
+}
+
+; CHECK-LABEL: test1
+; CHECK:   try
+; CHECK:     call      foo
+; --- Nested try/catch/end_try starts here
+; CHECK:     try
+; CHECK:       call      bar
+; CHECK:       call      bar
+; LOCAL:     catch     $push[[EXN0:[0-9]+]]=            # catch4:
+; LOCAL:       local.set  [[LOCAL:[0-9]+]], $pop[[EXN0]]
+; NOLOCAL:   catch     $[[REG:[0-9]+]]=
+; NOLOCAL:   local.copy  $[[DSTREG:[0-9]+]]=, $[[REG]]
+; CHECK:       br        1                               # 1: down to label4
+; CHECK:     end_try                                     # label5:
+; --- Nested try/catch/end_try ends here
+; CHECK:     return
+; CHECK:   catch     {{.*}}                              # catch3:
+; CHECK:     return
+; CHECK:   end_try                                       # label4:
+; LOCAL:   local.get  $push[[EXN1:[0-9]+]]=, [[LOCAL]]
+; LOCAL:   rethrow   $pop[[EXN1]]                        # to caller
+; NOLOCAL: rethrow   $[[DSTREG]]
+define void @test1() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+bb0:
+  invoke void @foo()
+          to label %bb1 unwind label %catch.dispatch0
+
+bb1:
+  call void @bar()
+  call void @bar()
+  ret void
+
+catch.dispatch0:                                  ; preds = %bb0
+  %0 = catchswitch within none [label %catch.start0] unwind to caller
+
+catch.start0:                                     ; preds = %catch.dispatch0
+  %1 = catchpad within %0 [i8* null]
+  %2 = call i8* @llvm.wasm.get.exception(token %1)
+  %3 = call i32 @llvm.wasm.get.ehselector(token %1)
+  catchret from %1 to label %try.cont
+
+try.cont:                                         ; preds = %bb0, %catch.start0
+  ret void
+}
+
+declare void @foo()
+declare void @bar()
+declare i32 @__gxx_wasm_personality_v0(...)
+declare i8* @llvm.wasm.get.exception(token)
+declare i32 @llvm.wasm.get.ehselector(token)
+declare i32 @llvm.eh.typeid.for(i8*)
+declare i8* @__cxa_begin_catch(i8*)
+declare void @__cxa_end_catch()
+declare void @__cxa_rethrow()
+declare void @__clang_call_terminate(i8*)
+declare void @_ZSt9terminatev()