[WebAssembly] Support Wasm EH + Wasm SjLj

D108960 added support for SjLj using Wasm EH instructions, which we call
Wasm SjLj going forward. (We call the old SjLj Emscripten SjLj) But it
did not support using Wasm EH and Wasm SjLj together. So far users of
Wasm EH had to use Wasm EH with Emscripten SjLj, which had a certain
limitation and it suffered from bigger code size increases as well.

This enables using Wasm EH and Wasm SjLj together.
1. This redirects `catchswitch` and `cleanupret` that unwind to caller
   to `catch.dispatch.longjmp` BB, which is a `catchswitch` BB that
   handles longjmps.
2. D108960 converted all longjmpable `call`s to `invokes` that unwind to
   `catch.dispatch.longjmp`. This CL checks if the `call` is embedded
   within another `catchpad`, and if so, makes it unwind to its nearest
   parent's unwind destination, rather than `catch.dispatch.longjmp`.
   This is necessary to preserve the scoping structure.

Reviewed By: dschuff

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

Author: aheejin

llvm/include/llvm/Transforms/Utils/Local.h

@@ -241,7 +241,7 @@ inline Align getKnownAlignment(Value *V, const DataLayout &DL,
 CallInst *createCallMatchingInvoke(InvokeInst *II);
 
 /// This function converts the specified invoek into a normall call.
-void changeToCall(InvokeInst *II, DomTreeUpdater *DTU = nullptr);
+CallInst *changeToCall(InvokeInst *II, DomTreeUpdater *DTU = nullptr);
 
 ///===---------------------------------------------------------------------===//
 ///  Dbg Intrinsic utilities

llvm/lib/Target/WebAssembly/WebAssemblyLowerEmscriptenEHSjLj.cpp

@@ -611,6 +611,8 @@ static bool canLongjmp(const Value *Callee) {
     return false;
   StringRef CalleeName = Callee->getName();
 
+  // TODO Include more functions or consider checking with mangled prefixes
+
   // The reason we include malloc/free here is to exclude the malloc/free
   // calls generated in setjmp prep / cleanup routines.
   if (CalleeName == "setjmp" || CalleeName == "malloc" || CalleeName == "free")
@@ -627,11 +629,50 @@ static bool canLongjmp(const Value *Callee) {
     return false;
 
   // Exception-catching related functions
-  if (CalleeName == "__cxa_begin_catch" || CalleeName == "__cxa_end_catch" ||
+  //
+  // We intentionally excluded __cxa_end_catch here even though it surely cannot
+  // longjmp, in order to maintain the unwind relationship from all existing
+  // catchpads (and calls within them) to catch.dispatch.longjmp.
+  //
+  // In Wasm EH + Wasm SjLj, we
+  // 1. Make all catchswitch and cleanuppad that unwind to caller unwind to
+  //    catch.dispatch.longjmp instead
+  // 2. Convert all longjmpable calls to invokes that unwind to
+  //    catch.dispatch.longjmp
+  // But catchswitch BBs are removed in isel, so if an EH catchswitch (generated
+  // from an exception)'s catchpad does not contain any calls that are converted
+  // into invokes unwinding to catch.dispatch.longjmp, this unwind relationship
+  // (EH catchswitch BB -> catch.dispatch.longjmp BB) is lost and
+  // catch.dispatch.longjmp BB can be placed before the EH catchswitch BB in
+  // CFGSort.
+  // int ret = setjmp(buf);
+  // try {
+  //   foo(); // longjmps
+  // } catch (...) {
+  // }
+  // Then in this code, if 'foo' longjmps, it first unwinds to 'catch (...)'
+  // catchswitch, and is not caught by that catchswitch because it is a longjmp,
+  // then it should next unwind to catch.dispatch.longjmp BB. But if this 'catch
+  // (...)' catchswitch -> catch.dispatch.longjmp unwind relationship is lost,
+  // it will not unwind to catch.dispatch.longjmp, producing an incorrect
+  // result.
+  //
+  // Every catchpad generated by Wasm C++ contains __cxa_end_catch, so we
+  // intentionally treat it as longjmpable to work around this problem. This is
+  // a hacky fix but an easy one.
+  //
+  // The comment block in findWasmUnwindDestinations() in
+  // SelectionDAGBuilder.cpp is addressing a similar problem.
+  if (CalleeName == "__cxa_begin_catch" ||
       CalleeName == "__cxa_allocate_exception" || CalleeName == "__cxa_throw" ||
       CalleeName == "__clang_call_terminate")
     return false;
 
+  // std::terminate, which is generated when another exception occurs while
+  // handling an exception, cannot longjmp.
+  if (CalleeName == "_ZSt9terminatev")
+    return false;
+
   // Otherwise we don't know
   return true;
 }
@@ -1271,16 +1312,19 @@ bool WebAssemblyLowerEmscriptenEHSjLj::runSjLjOnFunction(Function &F) {
   // Setjmp transformation
   SmallVector<PHINode *, 4> SetjmpRetPHIs;
   Function *SetjmpF = M.getFunction("setjmp");
-  for (User *U : SetjmpF->users()) {
-    auto *CI = dyn_cast<CallInst>(U);
-    // FIXME 'invoke' to setjmp can happen when we use Wasm EH + Wasm SjLj, but
-    // we don't support two being used together yet.
-    if (!CI)
-      report_fatal_error("Wasm EH + Wasm SjLj is not fully supported yet");
-    BasicBlock *BB = CI->getParent();
+  for (auto *U : make_early_inc_range(SetjmpF->users())) {
+    auto *CB = dyn_cast<CallBase>(U);
+    BasicBlock *BB = CB->getParent();
     if (BB->getParent() != &F) // in other function
       continue;
 
+    CallInst *CI = nullptr;
+    // setjmp cannot throw. So if it is an invoke, lower it to a call
+    if (auto *II = dyn_cast<InvokeInst>(CB))
+      CI = llvm::changeToCall(II);
+    else
+      CI = cast<CallInst>(CB);
+
     // The tail is everything right after the call, and will be reached once
     // when setjmp is called, and later when longjmp returns to the setjmp
     BasicBlock *Tail = SplitBlock(BB, CI->getNextNode());
@@ -1596,6 +1640,13 @@ void WebAssemblyLowerEmscriptenEHSjLj::handleLongjmpableCallsForEmscriptenSjLj(
     I->eraseFromParent();
 }
 
+static BasicBlock *getCleanupRetUnwindDest(const CleanupPadInst *CPI) {
+  for (const User *U : CPI->users())
+    if (const auto *CRI = dyn_cast<CleanupReturnInst>(U))
+      return CRI->getUnwindDest();
+  return nullptr;
+}
+
 // Create a catchpad in which we catch a longjmp's env and val arguments, test
 // if the longjmp corresponds to one of setjmps in the current function, and if
 // so, jump to the setjmp dispatch BB from which we go to one of post-setjmp
@@ -1747,15 +1798,66 @@ void WebAssemblyLowerEmscriptenEHSjLj::handleLongjmpableCallsForWasmSjLj(
       LongjmpableCalls.push_back(CI);
     }
   }
+
   for (auto *CI : LongjmpableCalls) {
     // Even if the callee function has attribute 'nounwind', which is true for
     // all C functions, it can longjmp, which means it can throw a Wasm
     // exception now.
     CI->removeFnAttr(Attribute::NoUnwind);
     if (Function *CalleeF = CI->getCalledFunction())
       CalleeF->removeFnAttr(Attribute::NoUnwind);
+
     // Change it to an invoke and make it unwind to the catch.dispatch.longjmp
-    // BB.
-    changeToInvokeAndSplitBasicBlock(CI, CatchDispatchLongjmpBB);
+    // BB. If the call is enclosed in another catchpad/cleanuppad scope, unwind
+    // to its parent pad's unwind destination instead to preserve the scope
+    // structure. It will eventually unwind to the catch.dispatch.longjmp.
+    SmallVector<OperandBundleDef, 1> Bundles;
+    BasicBlock *UnwindDest = nullptr;
+    if (auto Bundle = CI->getOperandBundle(LLVMContext::OB_funclet)) {
+      Instruction *FromPad = cast<Instruction>(Bundle->Inputs[0]);
+      while (!UnwindDest && FromPad) {
+        if (auto *CPI = dyn_cast<CatchPadInst>(FromPad)) {
+          UnwindDest = CPI->getCatchSwitch()->getUnwindDest();
+          FromPad = nullptr; // stop searching
+        } else if (auto *CPI = dyn_cast<CleanupPadInst>(FromPad)) {
+          // getCleanupRetUnwindDest() can return nullptr when
+          // 1. This cleanuppad's matching cleanupret uwninds to caller
+          // 2. There is no matching cleanupret because it ends with
+          //    unreachable.
+          // In case of 2, we need to traverse the parent pad chain.
+          UnwindDest = getCleanupRetUnwindDest(CPI);
+          FromPad = cast<Instruction>(CPI->getParentPad());
+        }
+      }
+    }
+    if (!UnwindDest)
+      UnwindDest = CatchDispatchLongjmpBB;
+    changeToInvokeAndSplitBasicBlock(CI, UnwindDest);
+  }
+
+  SmallVector<Instruction *, 16> ToErase;
+  for (auto &BB : F) {
+    if (auto *CSI = dyn_cast<CatchSwitchInst>(BB.getFirstNonPHI())) {
+      if (CSI != CatchSwitchLongjmp && CSI->unwindsToCaller()) {
+        IRB.SetInsertPoint(CSI);
+        ToErase.push_back(CSI);
+        auto *NewCSI = IRB.CreateCatchSwitch(CSI->getParentPad(),
+                                             CatchDispatchLongjmpBB, 1);
+        NewCSI->addHandler(*CSI->handler_begin());
+        NewCSI->takeName(CSI);
+        CSI->replaceAllUsesWith(NewCSI);
+      }
+    }
+
+    if (auto *CRI = dyn_cast<CleanupReturnInst>(BB.getTerminator())) {
+      if (CRI->unwindsToCaller()) {
+        IRB.SetInsertPoint(CRI);
+        ToErase.push_back(CRI);
+        IRB.CreateCleanupRet(CRI->getCleanupPad(), CatchDispatchLongjmpBB);
+      }
+    }
   }
+
+  for (Instruction *I : ToErase)
+    I->eraseFromParent();
 }

llvm/lib/Transforms/Utils/Local.cpp

@@ -2189,8 +2189,8 @@ CallInst *llvm::createCallMatchingInvoke(InvokeInst *II) {
   return NewCall;
 }
 
-/// changeToCall - Convert the specified invoke into a normal call.
-void llvm::changeToCall(InvokeInst *II, DomTreeUpdater *DTU) {
+// changeToCall - Convert the specified invoke into a normal call.
+CallInst *llvm::changeToCall(InvokeInst *II, DomTreeUpdater *DTU) {
   CallInst *NewCall = createCallMatchingInvoke(II);
   NewCall->takeName(II);
   NewCall->insertBefore(II);
@@ -2207,6 +2207,7 @@ void llvm::changeToCall(InvokeInst *II, DomTreeUpdater *DTU) {
   II->eraseFromParent();
   if (DTU)
     DTU->applyUpdates({{DominatorTree::Delete, BB, UnwindDestBB}});
+  return NewCall;
 }
 
 BasicBlock *llvm::changeToInvokeAndSplitBasicBlock(CallInst *CI,