[win][x64] Unwind v2 3/n: Add support for emitting unwind v2 information (equivalent to MSVC /d2epilogunwind) (#129142)

Adds support for emitting Windows x64 Unwind V2 information, includes
support `/d2epilogunwind` in clang-cl.

Unwind v2 adds information about the epilogs in functions such that the
unwinder can unwind even in the middle of an epilog, without having to
disassembly the function to see what has or has not been cleaned up.

Unwind v2 requires that all epilogs are in "canonical" form:
* If there was a stack allocation (fixed or dynamic) in the prolog, then
the first instruction in the epilog must be a stack deallocation.
* Next, for each `PUSH` in the prolog there must be a corresponding
`POP` instruction in exact reverse order.
* Finally, the epilog must end with the terminator.

This change adds a pass to validate epilogs in modules that have Unwind
v2 enabled and, if they pass, emits new pseudo instructions to MC that
1) note that the function is using unwind v2 and 2) mark the start of
the epilog (this is either the first `POP` if there is one, otherwise
the terminator instruction). If a function does not meet these
requirements, it is downgraded to Unwind v1 (i.e., these new pseudo
instructions are not emitted).

Note that the unwind v2 table only marks the size of the epilog in the
"header" unwind code, but it's possible for epilogs to use different
terminator instructions thus they are not all the same size. As a work
around for this, MC will assume that all terminator instructions are
1-byte long - this still works correctly with the Windows unwinder as it
is only using the size to do a range check to see if a thread is in an
epilog or not, and since the instruction pointer will never be in the
middle of an instruction and the terminator is always at the end of an
epilog the range check will function correctly. This does mean, however,
that the "at end" optimization (where an epilog unwind code can be
elided if the last epilog is at the end of the function) can only be
used if the terminator is 1-byte long.

One other complication with the implementation is that the unwind table
for a function is emitted during streaming, however we can't calculate
the distance between an epilog and the end of the function at that time
as layout hasn't been completed yet (thus some instructions may be
relaxed). To work around this, epilog unwind codes are emitted via a
fixup. This also means that we can't pre-emptively downgrade a function
to Unwind v1 if one of these offsets is too large, so instead we raise
an error (but I've passed through the location information, so the user
will know which of their functions is problematic).

1 files changed, 179 insertions, 14 deletions

diff --git a/llvm/lib/MC/MCWin64EH.cpp b/llvm/lib/MC/MCWin64EH.cpp
index e21f775..7fc8919 100644
--- a/llvm/lib/MC/MCWin64EH.cpp
+++ b/llvm/lib/MC/MCWin64EH.cpp
@@ -8,14 +8,57 @@
 
 #include "llvm/MC/MCWin64EH.h"
 #include "llvm/ADT/Twine.h"
+#include "llvm/MC/MCAssembler.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCObjectStreamer.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSymbol.h"
+#include "llvm/MC/MCValue.h"
 #include "llvm/Support/Win64EH.h"
+
 namespace llvm {
 class MCSection;
+
+/// MCExpr that represents the epilog unwind code in an unwind table.
+class MCUnwindV2EpilogTargetExpr final : public MCTargetExpr {
+  const MCSymbol *FunctionEnd;
+  const MCSymbol *UnwindV2Start;
+  const MCSymbol *EpilogEnd;
+  uint8_t EpilogSize;
+  SMLoc Loc;
+
+  MCUnwindV2EpilogTargetExpr(const WinEH::FrameInfo &FrameInfo,
+                             const WinEH::FrameInfo::Epilog &Epilog,
+                             uint8_t EpilogSize_)
+      : FunctionEnd(FrameInfo.FuncletOrFuncEnd),
+        UnwindV2Start(Epilog.UnwindV2Start), EpilogEnd(Epilog.End),
+        EpilogSize(EpilogSize_), Loc(Epilog.Loc) {}
+
+public:
+  static MCUnwindV2EpilogTargetExpr *
+  create(const WinEH::FrameInfo &FrameInfo,
+         const WinEH::FrameInfo::Epilog &Epilog, uint8_t EpilogSize_,
+         MCContext &Ctx) {
+    return new (Ctx) MCUnwindV2EpilogTargetExpr(FrameInfo, Epilog, EpilogSize_);
+  }
+
+  void printImpl(raw_ostream &OS, const MCAsmInfo *MAI) const override {
+    OS << ":epilog:";
+    UnwindV2Start->print(OS, MAI);
+  }
+
+  bool evaluateAsRelocatableImpl(MCValue &Res,
+                                 const MCAssembler *Asm) const override;
+
+  void visitUsedExpr(MCStreamer &Streamer) const override {
+    // Contains no sub-expressions.
+  }
+
+  MCFragment *findAssociatedFragment() const override {
+    return UnwindV2Start->getFragment();
+  }
+};
 }
 
 using namespace llvm;
@@ -163,20 +206,91 @@ static void EmitRuntimeFunction(MCStreamer &streamer,
                                              context), 4);
 }
 
+static std::optional<int64_t>
+GetOptionalAbsDifference(const MCAssembler &Assembler, const MCSymbol *LHS,
+                         const MCSymbol *RHS) {
+  MCContext &Context = Assembler.getContext();
+  const MCExpr *Diff =
+      MCBinaryExpr::createSub(MCSymbolRefExpr::create(LHS, Context),
+                              MCSymbolRefExpr::create(RHS, Context), Context);
+  // It should normally be possible to calculate the length of a function
+  // at this point, but it might not be possible in the presence of certain
+  // unusual constructs, like an inline asm with an alignment directive.
+  int64_t value;
+  if (!Diff->evaluateAsAbsolute(value, Assembler))
+    return std::nullopt;
+  return value;
+}
+
 static void EmitUnwindInfo(MCStreamer &streamer, WinEH::FrameInfo *info) {
   // If this UNWIND_INFO already has a symbol, it's already been emitted.
   if (info->Symbol)
     return;
 
   MCContext &context = streamer.getContext();
+  MCObjectStreamer *OS = (MCObjectStreamer *)(&streamer);
   MCSymbol *Label = context.createTempSymbol();
 
   streamer.emitValueToAlignment(Align(4));
   streamer.emitLabel(Label);
   info->Symbol = Label;
 
-  // Upper 3 bits are the version number (currently 1).
-  uint8_t flags = 0x01;
+  uint8_t numCodes = CountOfUnwindCodes(info->Instructions);
+  bool LastEpilogIsAtEnd = false;
+  bool AddPaddingEpilogCode = false;
+  uint8_t EpilogSize = 0;
+  bool EnableUnwindV2 = (info->Version >= 2) && !info->EpilogMap.empty();
+  if (EnableUnwindV2) {
+    auto &LastEpilog = info->EpilogMap.back().second;
+
+    // Calculate the size of the epilogs. Note that we +1 to the size so that
+    // the terminator instruction is also included in the epilog (the Windows
+    // unwinder does a simple range check versus the current instruction pointer
+    // so, although there are terminators that are large than 1 byte, the
+    // starting address of the terminator instruction will always be considered
+    // inside the epilog).
+    auto MaybeSize = GetOptionalAbsDifference(
+        OS->getAssembler(), LastEpilog.End, LastEpilog.UnwindV2Start);
+    if (!MaybeSize) {
+      context.reportError(LastEpilog.Loc,
+                          "Failed to evaluate epilog size for Unwind v2");
+      return;
+    }
+    assert(*MaybeSize >= 0);
+    if (*MaybeSize >= (int64_t)UINT8_MAX) {
+      context.reportError(LastEpilog.Loc,
+                          "Epilog size is too large for Unwind v2");
+      return;
+    }
+    EpilogSize = *MaybeSize + 1;
+
+    // If the last epilog is at the end of the function, we can use a special
+    // encoding for it. Because of our +1 trick for the size, this will only
+    // work where that final terminator instruction is 1 byte long.
+    auto LastEpilogToFuncEnd = GetOptionalAbsDifference(
+        OS->getAssembler(), info->FuncletOrFuncEnd, LastEpilog.UnwindV2Start);
+    LastEpilogIsAtEnd = (LastEpilogToFuncEnd == EpilogSize);
+
+    // If we have an odd number of epilog codes, we need to add a padding code.
+    size_t numEpilogCodes =
+        info->EpilogMap.size() + (LastEpilogIsAtEnd ? 0 : 1);
+    if ((numEpilogCodes % 2) != 0) {
+      AddPaddingEpilogCode = true;
+      numEpilogCodes++;
+    }
+
+    // Too many epilogs to handle.
+    if ((size_t)numCodes + numEpilogCodes > UINT8_MAX) {
+      context.reportError(info->FunctionLoc,
+                          "Too many unwind codes with Unwind v2 enabled");
+      return;
+    }
+
+    numCodes += numEpilogCodes;
+  }
+
+  // Upper 3 bits are the version number.
+  uint8_t flags = info->Version;
   if (info->ChainedParent)
     flags |= Win64EH::UNW_ChainInfo << 3;
   else {
@@ -192,7 +306,6 @@ static void EmitUnwindInfo(MCStreamer &streamer, WinEH::FrameInfo *info) {
   else
     streamer.emitInt8(0);
 
-  uint8_t numCodes = CountOfUnwindCodes(info->Instructions);
   streamer.emitInt8(numCodes);
 
   uint8_t frame = 0;
@@ -203,6 +316,35 @@ static void EmitUnwindInfo(MCStreamer &streamer, WinEH::FrameInfo *info) {
   }
   streamer.emitInt8(frame);
 
+  // Emit the epilog instructions.
+  if (EnableUnwindV2) {
+    MCDataFragment *DF = OS->getOrCreateDataFragment();
+
+    bool IsLast = true;
+    for (const auto &Epilog : llvm::reverse(info->EpilogMap)) {
+      if (IsLast) {
+        IsLast = false;
+        uint8_t Flags = LastEpilogIsAtEnd ? 0x01 : 0;
+        streamer.emitInt8(EpilogSize);
+        streamer.emitInt8((Flags << 4) | Win64EH::UOP_Epilog);
+
+        if (LastEpilogIsAtEnd)
+          continue;
+      }
+
+      // Each epilog is emitted as a fixup, since we can't measure the distance
+      // between the start of the epilog and the end of the function until
+      // layout has been completed.
+      auto *MCE = MCUnwindV2EpilogTargetExpr::create(*info, Epilog.second,
+                                                     EpilogSize, context);
+      MCFixup Fixup = MCFixup::create(DF->getContents().size(), MCE, FK_Data_2);
+      DF->getFixups().push_back(Fixup);
+      DF->appendContents(2, 0);
+    }
+  }
+  if (AddPaddingEpilogCode)
+    streamer.emitInt16(Win64EH::UOP_Epilog << 8);
+
   // Emit unwind instructions (in reverse order).
   uint8_t numInst = info->Instructions.size();
   for (uint8_t c = 0; c < numInst; ++c) {
@@ -234,6 +376,39 @@ static void EmitUnwindInfo(MCStreamer &streamer, WinEH::FrameInfo *info) {
   }
 }
 
+bool MCUnwindV2EpilogTargetExpr::evaluateAsRelocatableImpl(
+    MCValue &Res, const MCAssembler *Asm) const {
+  // Calculate the offset to this epilog, and validate it's within the allowed
+  // range.
+  auto Offset = GetOptionalAbsDifference(*Asm, FunctionEnd, UnwindV2Start);
+  if (!Offset) {
+    Asm->getContext().reportError(
+        Loc, "Failed to evaluate epilog offset for Unwind v2");
+    return false;
+  }
+  assert(*Offset > 0);
+  constexpr uint16_t MaxEpilogOffset = 0x0fff;
+  if (*Offset > MaxEpilogOffset) {
+    Asm->getContext().reportError(Loc,
+                                  "Epilog offset is too large for Unwind v2");
+    return false;
+  }
+
+  // Sanity check that all epilogs are the same size.
+  auto Size = GetOptionalAbsDifference(*Asm, EpilogEnd, UnwindV2Start);
+  if (Size != (EpilogSize - 1)) {
+    Asm->getContext().reportError(
+        Loc,
+        "Size of this epilog does not match size of last epilog in function");
+    return false;
+  }
+
+  auto HighBits = *Offset >> 8;
+  Res = MCValue::get((HighBits << 12) | (Win64EH::UOP_Epilog << 8) |
+                     (*Offset & 0xFF));
+  return true;
+}
+
 void llvm::Win64EH::UnwindEmitter::Emit(MCStreamer &Streamer) const {
   // Emit the unwind info structs first.
   for (const auto &CFI : Streamer.getWinFrameInfos()) {
@@ -276,18 +451,8 @@ static const MCExpr *GetSubDivExpr(MCStreamer &Streamer, const MCSymbol *LHS,
 static std::optional<int64_t> GetOptionalAbsDifference(MCStreamer &Streamer,
                                                        const MCSymbol *LHS,
                                                        const MCSymbol *RHS) {
-  MCContext &Context = Streamer.getContext();
-  const MCExpr *Diff =
-      MCBinaryExpr::createSub(MCSymbolRefExpr::create(LHS, Context),
-                              MCSymbolRefExpr::create(RHS, Context), Context);
   MCObjectStreamer *OS = (MCObjectStreamer *)(&Streamer);
-  // It should normally be possible to calculate the length of a function
-  // at this point, but it might not be possible in the presence of certain
-  // unusual constructs, like an inline asm with an alignment directive.
-  int64_t value;
-  if (!Diff->evaluateAsAbsolute(value, OS->getAssembler()))
-    return std::nullopt;
-  return value;
+  return GetOptionalAbsDifference(OS->getAssembler(), LHS, RHS);
 }
 
 static int64_t GetAbsDifference(MCStreamer &Streamer, const MCSymbol *LHS,