[WebAssembly] Ensure terminate pads are a single BB

This ensures every single terminate pad is a single BB in the form of:
```
%exn = catch $__cpp_exception
call @__clang_call_terminate(%exn)
unreachable
```

This is a preparation for HandleEHTerminatePads pass, which will be
added in a later CL and will run after CFGStackify. That pass duplicates
terminate pads with a `catch_all` instruction, and duplicating it
becomes simpler if we can ensure every terminate pad is a single BB.

Reviewed By: dschuff, tlively

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

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

@@ -178,27 +178,6 @@ getLatestInsertPos(MachineBasicBlock *MBB,
   return InsertPos;
 }
 
-// Find a catch instruction and its destination register within an EH pad.
-static MachineInstr *findCatch(MachineBasicBlock *EHPad, Register &ExnReg) {
-  assert(EHPad->isEHPad());
-  MachineInstr *Catch = nullptr;
-  for (auto &MI : *EHPad) {
-    if (WebAssembly::isCatch(MI.getOpcode())) {
-      Catch = &MI;
-      ExnReg = Catch->getOperand(0).getReg();
-      break;
-    }
-  }
-  assert(Catch && "EH pad does not have a catch");
-  assert(ExnReg != 0 && "Invalid register");
-  return Catch;
-}
-
-static MachineInstr *findCatch(MachineBasicBlock *EHPad) {
-  Register Dummy;
-  return findCatch(EHPad, Dummy);
-}
-
 void WebAssemblyCFGStackify::registerScope(MachineInstr *Begin,
                                            MachineInstr *End) {
   BeginToEnd[Begin] = End;
@@ -874,7 +853,10 @@ void WebAssemblyCFGStackify::fixEndsAtEndOfFunction(MachineFunction &MF) {
         // instructions before its corresponding 'catch' too.
         auto *EHPad = TryToEHPad.lookup(EndToBegin[&MI]);
         assert(EHPad);
-        Worklist.push_back(std::next(findCatch(EHPad)->getReverseIterator()));
+        auto NextIt =
+            std::next(WebAssembly::findCatch(EHPad)->getReverseIterator());
+        if (NextIt != EHPad->rend())
+          Worklist.push_back(NextIt);
         LLVM_FALLTHROUGH;
       }
       case WebAssembly::END_BLOCK:

llvm/lib/Target/WebAssembly/WebAssemblyLateEHPrepare.cpp

@@ -38,6 +38,7 @@ class WebAssemblyLateEHPrepare final : public MachineFunctionPass {
   bool addCatchAlls(MachineFunction &MF);
   bool replaceFuncletReturns(MachineFunction &MF);
   bool removeUnnecessaryUnreachables(MachineFunction &MF);
+  bool ensureSingleBBTermPads(MachineFunction &MF);
   bool restoreStackPointer(MachineFunction &MF);
 
   MachineBasicBlock *getMatchingEHPad(MachineInstr *MI);
@@ -127,6 +128,7 @@ bool WebAssemblyLateEHPrepare::runOnMachineFunction(MachineFunction &MF) {
     Changed |= hoistCatches(MF);
     Changed |= addCatchAlls(MF);
     Changed |= replaceFuncletReturns(MF);
+    Changed |= ensureSingleBBTermPads(MF);
   }
   Changed |= removeUnnecessaryUnreachables(MF);
   if (MF.getFunction().hasPersonalityFn())
@@ -285,6 +287,81 @@ bool WebAssemblyLateEHPrepare::removeUnnecessaryUnreachables(
   return Changed;
 }
 
+// Clang-generated terminate pads are an single-BB EH pad in the form of
+// termpad:
+//   %exn = catch $__cpp_exception
+//   call @__clang_call_terminate(%exn)
+//   unreachable
+// (There can be local.set and local.gets before the call if we didn't run
+// RegStackify)
+// But code transformations can change or add more control flow, so the call to
+// __clang_call_terminate() function may not be in the original EH pad anymore.
+// This ensures every terminate pad is a single BB in the form illustrated
+// above.
+//
+// This is preparation work for the HandleEHTerminatePads pass later, which
+// duplicates terminate pads both for 'catch' and 'catch_all'. Refer to
+// WebAssemblyHandleEHTerminatePads.cpp for details.
+bool WebAssemblyLateEHPrepare::ensureSingleBBTermPads(MachineFunction &MF) {
+  const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
+
+  // Find calls to __clang_call_terminate()
+  SmallVector<MachineInstr *, 8> ClangCallTerminateCalls;
+  SmallPtrSet<MachineBasicBlock *, 8> TermPads;
+  for (auto &MBB : MF) {
+    for (auto &MI : MBB) {
+      if (MI.isCall()) {
+        const MachineOperand &CalleeOp = MI.getOperand(0);
+        if (CalleeOp.isGlobal() && CalleeOp.getGlobal()->getName() ==
+                                       WebAssembly::ClangCallTerminateFn) {
+          MachineBasicBlock *EHPad = getMatchingEHPad(&MI);
+          assert(EHPad && "No matching EH pad for __clang_call_terminate");
+          // In case a __clang_call_terminate call is duplicated during code
+          // transformation so one terminate pad contains multiple
+          // __clang_call_terminate calls, we only count one of them
+          if (TermPads.insert(EHPad).second)
+            ClangCallTerminateCalls.push_back(&MI);
+        }
+      }
+    }
+  }
+
+  bool Changed = false;
+  for (auto *Call : ClangCallTerminateCalls) {
+    MachineBasicBlock *EHPad = getMatchingEHPad(Call);
+    assert(EHPad && "No matching EH pad for __clang_call_terminate");
+
+    // If it is already the form we want, skip it
+    if (Call->getParent() == EHPad &&
+        Call->getNextNode()->getOpcode() == WebAssembly::UNREACHABLE)
+      continue;
+
+    // In case the __clang_call_terminate() call is not in its matching EH pad,
+    // move the call to the end of EH pad and add an unreachable instruction
+    // after that. Delete all successors and their children if any, because here
+    // the program terminates.
+    Changed = true;
+    // This runs after hoistCatches(), so catch instruction should be at the top
+    MachineInstr *Catch = WebAssembly::findCatch(EHPad);
+    assert(Catch && "EH pad does not have a catch instruction");
+    // Takes the result register of the catch instruction as argument. There may
+    // have been some other local.set/local.gets in between, but at this point
+    // we don't care.
+    Call->getOperand(1).setReg(Catch->getOperand(0).getReg());
+    auto InsertPos = std::next(MachineBasicBlock::iterator(Catch));
+    EHPad->insert(InsertPos, Call->removeFromParent());
+    BuildMI(*EHPad, InsertPos, Call->getDebugLoc(),
+            TII.get(WebAssembly::UNREACHABLE));
+    EHPad->erase(InsertPos, EHPad->end());
+    SmallVector<MachineBasicBlock *, 8> Succs(EHPad->succ_begin(),
+                                              EHPad->succ_end());
+    for (auto *Succ : Succs)
+      EHPad->removeSuccessor(Succ);
+    eraseDeadBBsAndChildren(Succs);
+  }
+  return Changed;
+}
+
 // After the stack is unwound due to a thrown exception, the __stack_pointer
 // global can point to an invalid address. This inserts instructions that
 // restore __stack_pointer global.

llvm/lib/Target/WebAssembly/WebAssemblyUtilities.cpp

@@ -115,3 +115,17 @@ WebAssembly::getOrCreateFunctionTableSymbol(MCContext &Ctx,
   }
   return Sym;
 }
+
+// Find a catch instruction from an EH pad.
+MachineInstr *WebAssembly::findCatch(MachineBasicBlock *EHPad) {
+  assert(EHPad->isEHPad());
+  auto Pos = EHPad->begin();
+  // Skip any label or debug instructions. Also skip 'end' marker instructions
+  // that may exist after marker placement in CFGStackify.
+  while (Pos != EHPad->end() &&
+         (Pos->isLabel() || Pos->isDebugInstr() || isMarker(Pos->getOpcode())))
+    Pos++;
+  if (Pos != EHPad->end() && WebAssembly::isCatch(Pos->getOpcode()))
+    return &*Pos;
+  return nullptr;
+}

llvm/lib/Target/WebAssembly/WebAssemblyUtilities.h

@@ -17,6 +17,7 @@
 
 namespace llvm {
 
+class MachineBasicBlock;
 class MachineInstr;
 class MachineOperand;
 class MCContext;
@@ -45,6 +46,10 @@ const MachineOperand &getCalleeOp(const MachineInstr &MI);
 MCSymbolWasm *getOrCreateFunctionTableSymbol(MCContext &Ctx,
                                              const StringRef &Name);
 
+/// Find a catch instruction from an EH pad. Returns null if no catch
+/// instruction found or the catch is in an invalid location.
+MachineInstr *findCatch(MachineBasicBlock *EHPad);
+
 } // end namespace WebAssembly
 
 } // end namespace llvm

llvm/test/CodeGen/WebAssembly/exception.ll

@@ -1,4 +1,5 @@
 ; RUN: llc < %s -asm-verbose=false -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals -wasm-keep-registers -exception-model=wasm -mattr=+exception-handling -verify-machineinstrs | FileCheck -allow-deprecated-dag-overlap %s
+; RUN: llc < %s -asm-verbose=false -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals -wasm-keep-registers -exception-model=wasm -mattr=+exception-handling -verify-machineinstrs -O0 | FileCheck -allow-deprecated-dag-overlap --check-prefix=NOOPT %s
 ; RUN: llc < %s -disable-wasm-fallthrough-return-opt -wasm-keep-registers -exception-model=wasm -mattr=+exception-handling
 
 target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
@@ -174,6 +175,73 @@ terminate:                                        ; preds = %ehcleanup
   unreachable
 }
 
+; Tests a case when there are multiple BBs within a terminate pad. This kind of
+; structure is not generated by clang, but can generated by code
+; transformations. After LateEHPrepare, there should be a single 'terminate' BB
+; with these instructions:
+
+; %exn = catch $__cpp_exception
+; call @__clang_call_terminate(%exn)
+; unreachable
+
+; NOOPT-LABEL: test_split_terminatepad
+define void @test_split_terminatepad(i1 %arg) personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+entry:
+  invoke void @foo()
+          to label %try.cont unwind label %catch.dispatch
+
+catch.dispatch:                                   ; preds = %entry
+  %0 = catchswitch within none [label %catch.start] unwind to caller
+
+; NOOPT:      catch
+catch.start:                                      ; preds = %catch.dispatch
+  %1 = catchpad within %0 [i8* null]
+  %2 = call i8* @llvm.wasm.get.exception(token %1)
+  %3 = call i32 @llvm.wasm.get.ehselector(token %1)
+  %4 = call i8* @__cxa_begin_catch(i8* %2) [ "funclet"(token %1) ]
+  invoke void @foo() [ "funclet"(token %1) ]
+          to label %invoke.cont1 unwind label %ehcleanup
+
+invoke.cont1:                                     ; preds = %catch.start
+  call void @__cxa_end_catch() [ "funclet"(token %1) ]
+  catchret from %1 to label %try.cont
+
+try.cont:                                         ; preds = %invoke.cont1, %entry
+  ret void
+
+; NOOPT:      catch_all
+ehcleanup:                                        ; preds = %catch.start
+  %5 = cleanuppad within %1 []
+  invoke void @__cxa_end_catch() [ "funclet"(token %5) ]
+          to label %invoke.cont2 unwind label %terminate
+
+invoke.cont2:                                     ; preds = %ehcleanup
+  cleanupret from %5 unwind to caller
+
+; This weird structure of split terminate pads are not generated by clang, but
+; we cannot guarantee this kind of multi-BB terminate pads cannot be generated
+; by code transformations. This structure is manually created for this test.
+; NOOPT:      catch     $[[EXN:[0-9]+]]=, __cpp_exception
+; NOOPT-NEXT: global.set  __stack_pointer
+; NOOPT-NEXT: call  __clang_call_terminate, $[[EXN]]
+; NOOPT-NEXT: unreachable
+
+terminate:                                        ; preds = %ehcleanup
+  %6 = cleanuppad within %5 []
+  %7 = call i8* @llvm.wasm.get.exception(token %6)
+  br i1 %arg, label %terminate.split1, label %terminate.split2
+
+terminate.split1:
+  call void @__clang_call_terminate(i8* %7) [ "funclet"(token %6) ]
+  unreachable
+
+terminate.split2:
+  ; This is to test a hypothetical case that a call to __clang_call_terminate is
+  ; duplicated within a terminate pad
+  call void @__clang_call_terminate(i8* %7) [ "funclet"(token %6) ]
+  unreachable
+}
+
 ; Tests prologues and epilogues are not generated within EH scopes.
 ; They should not be treated as funclets; BBs starting with a catch instruction
 ; should not have a prologue, and BBs ending with a catchret/cleanupret should