4 files changed, 125 insertions, 11 deletions

diff --git a/llvm/lib/Target/X86/CMakeLists.txt b/llvm/lib/Target/X86/CMakeLists.txt
index ed34a59df4a..58f2292dd4c 100644
--- a/llvm/lib/Target/X86/CMakeLists.txt
+++ b/llvm/lib/Target/X86/CMakeLists.txt
@@ -23,6 +23,7 @@ add_public_tablegen_target(X86CommonTableGen)
 
 set(sources
   X86AsmPrinter.cpp
+  X86AvoidTrailingCall.cpp
   X86CallFrameOptimization.cpp
   X86CallingConv.cpp
   X86CallLowering.cpp
diff --git a/llvm/lib/Target/X86/X86.h b/llvm/lib/Target/X86/X86.h
index 573888a57fe..6840fc12751 100644
--- a/llvm/lib/Target/X86/X86.h
+++ b/llvm/lib/Target/X86/X86.h
@@ -81,6 +81,12 @@ FunctionPass *createX86FlagsCopyLoweringPass();
 /// Return a pass that expands WinAlloca pseudo-instructions.
 FunctionPass *createX86WinAllocaExpander();
 
+/// Return a pass that inserts int3 at the end of the function if it ends with a
+/// CALL instruction. The pass does the same for each funclet as well. This
+/// ensures that the open interval of function start and end PCs contains all
+/// return addresses for the benefit of the Windows x64 unwinder.
+FunctionPass *createX86AvoidTrailingCallPass();
+
 /// Return a pass that optimizes the code-size of x86 call sequences. This is
 /// done by replacing esp-relative movs with pushes.
 FunctionPass *createX86CallFrameOptimization();
diff --git a/llvm/lib/Target/X86/X86AvoidTrailingCall.cpp b/llvm/lib/Target/X86/X86AvoidTrailingCall.cpp
new file mode 100644
index 00000000000..fb4f9e2901d
--- /dev/null
+++ b/llvm/lib/Target/X86/X86AvoidTrailingCall.cpp
@@ -0,0 +1,108 @@
+//===----- X86AvoidTrailingCall.cpp - Insert int3 after trailing calls ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// The Windows x64 unwinder has trouble unwinding the stack when a return
+// address points to the end of the function. This pass maintains the invariant
+// that every return address is inside the bounds of its parent function or
+// funclet by inserting int3 if the last instruction would otherwise be a call.
+//
+//===----------------------------------------------------------------------===//
+
+#include "X86.h"
+#include "X86InstrInfo.h"
+#include "X86Subtarget.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+
+#define DEBUG_TYPE "x86-avoid-trailing-call"
+
+using namespace llvm;
+
+namespace {
+
+class X86AvoidTrailingCallPass : public MachineFunctionPass {
+public:
+  X86AvoidTrailingCallPass() : MachineFunctionPass(ID) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+private:
+  StringRef getPassName() const override {
+    return "X86 avoid trailing call pass";
+  }
+  static char ID;
+};
+
+char X86AvoidTrailingCallPass::ID = 0;
+
+} // end anonymous namespace
+
+FunctionPass *llvm::createX86AvoidTrailingCallPass() {
+  return new X86AvoidTrailingCallPass();
+}
+
+// A real instruction is a non-meta, non-pseudo instruction.  Some pseudos
+// expand to nothing, and some expand to code. This logic conservatively assumes
+// they might expand to nothing.
+static bool isRealInstruction(MachineInstr &MI) {
+  return !MI.isPseudo() && !MI.isMetaInstruction();
+}
+
+// Return true if this is a call instruction, but not a tail call.
+static bool isCallInstruction(const MachineInstr &MI) {
+  return MI.isCall() && !MI.isReturn();
+}
+
+bool X86AvoidTrailingCallPass::runOnMachineFunction(MachineFunction &MF) {
+  const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
+  const X86InstrInfo &TII = *STI.getInstrInfo();
+  assert(STI.isTargetWin64() && "pass only runs on Win64");
+
+  // FIXME: Perhaps this pass should also replace SEH_Epilogue by inserting nops
+  // before epilogues.
+
+  bool Changed = false;
+  for (MachineBasicBlock &MBB : MF) {
+    // Look for basic blocks that precede funclet entries or are at the end of
+    // the function.
+    MachineBasicBlock *NextMBB = MBB.getNextNode();
+    if (NextMBB && !NextMBB->isEHFuncletEntry())
+      continue;
+
+    // Find the last real instruction in this block, or previous blocks if this
+    // block is empty.
+    MachineBasicBlock::reverse_iterator LastRealInstr;
+    for (MachineBasicBlock &RMBB :
+         make_range(MBB.getReverseIterator(), MF.rend())) {
+      LastRealInstr = llvm::find_if(reverse(RMBB), isRealInstruction);
+      if (LastRealInstr != RMBB.rend())
+        break;
+    }
+
+    // Do nothing if this function or funclet has no instructions.
+    if (LastRealInstr == MF.begin()->rend())
+      continue;
+
+    // If this is a call instruction, insert int3 right after it with the same
+    // DebugLoc. Convert back to a forward iterator and advance the insertion
+    // position once.
+    if (isCallInstruction(*LastRealInstr)) {
+      LLVM_DEBUG({
+        dbgs() << "inserting int3 after trailing call instruction:\n";
+        LastRealInstr->dump();
+        dbgs() << '\n';
+      });
+
+      MachineBasicBlock::iterator MBBI = std::next(LastRealInstr.getReverse());
+      BuildMI(*LastRealInstr->getParent(), MBBI, LastRealInstr->getDebugLoc(),
+              TII.get(X86::INT3));
+      Changed = true;
+    }
+  }
+
+  return Changed;
+}
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index d33a7a0b963..f6184a3c1f1 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -219,17 +219,9 @@ X86TargetMachine::X86TargetMachine(const Target &T, const Triple &TT,
           getEffectiveX86CodeModel(CM, JIT, TT.getArch() == Triple::x86_64),
           OL),
       TLOF(createTLOF(getTargetTriple())) {
-  // Windows stack unwinder gets confused when execution flow "falls through"
-  // after a call to 'noreturn' function.
-  // To prevent that, we emit a trap for 'unreachable' IR instructions.
-  // (which on X86, happens to be the 'ud2' instruction)
   // On PS4, the "return address" of a 'noreturn' call must still be within
   // the calling function, and TrapUnreachable is an easy way to get that.
-  // The check here for 64-bit windows is a bit icky, but as we're unlikely
-  // to ever want to mix 32 and 64-bit windows code in a single module
-  // this should be fine.
-  if ((TT.isOSWindows() && TT.getArch() == Triple::x86_64) || TT.isPS4() ||
-      TT.isOSBinFormatMachO()) {
+  if (TT.isPS4() || TT.isOSBinFormatMachO()) {
     this->Options.TrapUnreachable = true;
     this->Options.NoTrapAfterNoreturn = TT.isOSBinFormatMachO();
   }
@@ -518,12 +510,19 @@ void X86PassConfig::addPreEmitPass() {
 }
 
 void X86PassConfig::addPreEmitPass2() {
+  const Triple &TT = TM->getTargetTriple();
+  const MCAsmInfo *MAI = TM->getMCAsmInfo();
+
   addPass(createX86RetpolineThunksPass());
+
+  // Insert extra int3 instructions after trailing call instructions to avoid
+  // issues in the unwinder.
+  if (TT.isOSWindows() && TT.getArch() == Triple::x86_64)
+    addPass(createX86AvoidTrailingCallPass());
+
   // Verify basic block incoming and outgoing cfa offset and register values and
   // correct CFA calculation rule where needed by inserting appropriate CFI
   // instructions.
-  const Triple &TT = TM->getTargetTriple();
-  const MCAsmInfo *MAI = TM->getMCAsmInfo();
   if (!TT.isOSDarwin() &&
       (!TT.isOSWindows() ||
        MAI->getExceptionHandlingType() == ExceptionHandling::DwarfCFI))