[win][x64] Unwind v2 4/4: Use chaining to split unwind info if needed

Author: dpaoliello

llvm/include/llvm/MC/MCStreamer.h

@@ -1024,8 +1024,7 @@ class LLVM_ABI MCStreamer {
   /// for the frame.  We cannot use the End marker, as it is not set at the
   /// point of emitting .xdata, in order to indicate that the frame is active.
   virtual void emitWinCFIFuncletOrFuncEnd(SMLoc Loc = SMLoc());
-  virtual void emitWinCFIStartChained(SMLoc Loc = SMLoc());
-  virtual void emitWinCFIEndChained(SMLoc Loc = SMLoc());
+  virtual void emitWinCFISplitChained(SMLoc Loc = SMLoc());
   virtual void emitWinCFIPushReg(MCRegister Register, SMLoc Loc = SMLoc());
   virtual void emitWinCFISetFrame(MCRegister Register, unsigned Offset,
                                   SMLoc Loc = SMLoc());

llvm/include/llvm/MC/MCWinEH.h

@@ -59,7 +59,7 @@ struct FrameInfo {
   uint8_t Version = DefaultVersion;
 
   int LastFrameInst = -1;
-  const FrameInfo *ChainedParent = nullptr;
+  FrameInfo *ChainedParent = nullptr;
   std::vector<Instruction> Instructions;
   struct Epilog {
     std::vector<Instruction> Instructions;
@@ -90,9 +90,9 @@ struct FrameInfo {
   FrameInfo(const MCSymbol *Function, const MCSymbol *BeginFuncEHLabel)
       : Begin(BeginFuncEHLabel), Function(Function) {}
   FrameInfo(const MCSymbol *Function, const MCSymbol *BeginFuncEHLabel,
-            const FrameInfo *ChainedParent)
+            FrameInfo *ChainedParent)
       : Begin(BeginFuncEHLabel), Function(Function),
-        ChainedParent(ChainedParent) {}
+        Version(ChainedParent->Version), ChainedParent(ChainedParent) {}
 
   bool empty() const {
     if (!Instructions.empty())

llvm/lib/MC/MCAsmStreamer.cpp

@@ -375,8 +375,7 @@ class MCAsmStreamer final : public MCStreamer {
   void emitWinCFIStartProc(const MCSymbol *Symbol, SMLoc Loc) override;
   void emitWinCFIEndProc(SMLoc Loc) override;
   void emitWinCFIFuncletOrFuncEnd(SMLoc Loc) override;
-  void emitWinCFIStartChained(SMLoc Loc) override;
-  void emitWinCFIEndChained(SMLoc Loc) override;
+  void emitWinCFISplitChained(SMLoc Loc) override;
   void emitWinCFIPushReg(MCRegister Register, SMLoc Loc) override;
   void emitWinCFISetFrame(MCRegister Register, unsigned Offset,
                           SMLoc Loc) override;
@@ -2175,17 +2174,10 @@ void MCAsmStreamer::emitWinCFIFuncletOrFuncEnd(SMLoc Loc) {
   EmitEOL();
 }
 
-void MCAsmStreamer::emitWinCFIStartChained(SMLoc Loc) {
-  MCStreamer::emitWinCFIStartChained(Loc);
+void MCAsmStreamer::emitWinCFISplitChained(SMLoc Loc) {
+  MCStreamer::emitWinCFISplitChained(Loc);
 
-  OS << "\t.seh_startchained";
-  EmitEOL();
-}
-
-void MCAsmStreamer::emitWinCFIEndChained(SMLoc Loc) {
-  MCStreamer::emitWinCFIEndChained(Loc);
-
-  OS << "\t.seh_endchained";
+  OS << "\t.seh_splitchained";
   EmitEOL();
 }
 

llvm/lib/MC/MCParser/COFFAsmParser.cpp

@@ -80,10 +80,8 @@ class COFFAsmParser : public MCAsmParserExtension {
         ".seh_endproc");
     addDirectiveHandler<&COFFAsmParser::parseSEHDirectiveEndFuncletOrFunc>(
         ".seh_endfunclet");
-    addDirectiveHandler<&COFFAsmParser::parseSEHDirectiveStartChained>(
-        ".seh_startchained");
-    addDirectiveHandler<&COFFAsmParser::parseSEHDirectiveEndChained>(
-        ".seh_endchained");
+    addDirectiveHandler<&COFFAsmParser::parseSEHDirectiveSplitChained>(
+        ".seh_splitchained");
     addDirectiveHandler<&COFFAsmParser::parseSEHDirectiveHandler>(
         ".seh_handler");
     addDirectiveHandler<&COFFAsmParser::parseSEHDirectiveHandlerData>(
@@ -143,8 +141,7 @@ class COFFAsmParser : public MCAsmParserExtension {
   bool parseSEHDirectiveStartProc(StringRef, SMLoc);
   bool parseSEHDirectiveEndProc(StringRef, SMLoc);
   bool parseSEHDirectiveEndFuncletOrFunc(StringRef, SMLoc);
-  bool parseSEHDirectiveStartChained(StringRef, SMLoc);
-  bool parseSEHDirectiveEndChained(StringRef, SMLoc);
+  bool parseSEHDirectiveSplitChained(StringRef, SMLoc);
   bool parseSEHDirectiveHandler(StringRef, SMLoc);
   bool parseSEHDirectiveHandlerData(StringRef, SMLoc);
   bool parseSEHDirectiveAllocStack(StringRef, SMLoc);
@@ -685,15 +682,9 @@ bool COFFAsmParser::parseSEHDirectiveEndFuncletOrFunc(StringRef, SMLoc Loc) {
   return false;
 }
 
-bool COFFAsmParser::parseSEHDirectiveStartChained(StringRef, SMLoc Loc) {
+bool COFFAsmParser::parseSEHDirectiveSplitChained(StringRef, SMLoc Loc) {
   Lex();
-  getStreamer().emitWinCFIStartChained(Loc);
-  return false;
-}
-
-bool COFFAsmParser::parseSEHDirectiveEndChained(StringRef, SMLoc Loc) {
-  Lex();
-  getStreamer().emitWinCFIEndChained(Loc);
+  getStreamer().emitWinCFISplitChained(Loc);
   return false;
 }
 

llvm/lib/MC/MCStreamer.cpp

@@ -749,13 +749,14 @@ void MCStreamer::emitWinCFIEndProc(SMLoc Loc) {
   WinEH::FrameInfo *CurFrame = EnsureValidWinFrameInfo(Loc);
   if (!CurFrame)
     return;
-  if (CurFrame->ChainedParent)
-    getContext().reportError(Loc, "Not all chained regions terminated!");
 
   MCSymbol *Label = emitCFILabel();
   CurFrame->End = Label;
-  if (!CurFrame->FuncletOrFuncEnd)
-    CurFrame->FuncletOrFuncEnd = CurFrame->End;
+  const MCSymbol **FuncletOrFuncEndPtr =
+      CurFrame->ChainedParent ? &CurFrame->ChainedParent->FuncletOrFuncEnd
+                              : &CurFrame->FuncletOrFuncEnd;
+  if (!*FuncletOrFuncEndPtr)
+    *FuncletOrFuncEndPtr = CurFrame->End;
 
   for (size_t I = CurrentProcWinFrameInfoStartIndex, E = WinFrameInfos.size();
        I != E; ++I)
@@ -767,48 +768,49 @@ void MCStreamer::emitWinCFIFuncletOrFuncEnd(SMLoc Loc) {
   WinEH::FrameInfo *CurFrame = EnsureValidWinFrameInfo(Loc);
   if (!CurFrame)
     return;
-  if (CurFrame->ChainedParent)
-    getContext().reportError(Loc, "Not all chained regions terminated!");
 
   MCSymbol *Label = emitCFILabel();
-  CurFrame->FuncletOrFuncEnd = Label;
+  const MCSymbol **FuncletOrFuncEndPtr =
+      CurFrame->ChainedParent ? &CurFrame->ChainedParent->FuncletOrFuncEnd
+                              : &CurFrame->FuncletOrFuncEnd;
+  *FuncletOrFuncEndPtr = Label;
 }
 
-void MCStreamer::emitWinCFIStartChained(SMLoc Loc) {
+void MCStreamer::emitWinCFISplitChained(SMLoc Loc) {
   WinEH::FrameInfo *CurFrame = EnsureValidWinFrameInfo(Loc);
   if (!CurFrame)
     return;
 
-  MCSymbol *StartProc = emitCFILabel();
-
-  WinFrameInfos.emplace_back(std::make_unique<WinEH::FrameInfo>(
-      CurFrame->Function, StartProc, CurFrame));
-  CurrentWinFrameInfo = WinFrameInfos.back().get();
-  CurrentWinFrameInfo->TextSection = getCurrentSectionOnly();
-}
-
-void MCStreamer::emitWinCFIEndChained(SMLoc Loc) {
-  WinEH::FrameInfo *CurFrame = EnsureValidWinFrameInfo(Loc);
-  if (!CurFrame)
-    return;
-  if (!CurFrame->ChainedParent)
+  if (!CurFrame->PrologEnd && !CurFrame->ChainedParent)
     return getContext().reportError(
-        Loc, "End of a chained region outside a chained region!");
+        Loc, "can't split into a new chained region (.seh_splitchained) in the "
+             "middle of a prolog in " +
+                 CurFrame->Function->getName());
 
   MCSymbol *Label = emitCFILabel();
 
+  // Complete the current frame before starting a new, chained one.
   CurFrame->End = Label;
-  CurrentWinFrameInfo = const_cast<WinEH::FrameInfo *>(CurFrame->ChainedParent);
+
+  // All chained frames point to the same parent.
+  WinEH::FrameInfo *ChainedParent =
+      CurFrame->ChainedParent ? CurFrame->ChainedParent : CurFrame;
+
+  WinFrameInfos.emplace_back(std::make_unique<WinEH::FrameInfo>(
+      CurFrame->Function, Label, ChainedParent));
+  CurrentWinFrameInfo = WinFrameInfos.back().get();
+  CurrentWinFrameInfo->TextSection = getCurrentSectionOnly();
 }
 
 void MCStreamer::emitWinEHHandler(const MCSymbol *Sym, bool Unwind, bool Except,
                                   SMLoc Loc) {
   WinEH::FrameInfo *CurFrame = EnsureValidWinFrameInfo(Loc);
   if (!CurFrame)
     return;
-  if (CurFrame->ChainedParent)
-    return getContext().reportError(
-        Loc, "Chained unwind areas can't have handlers!");
+
+  // Handlers are always associated with the parent frame.
+  CurFrame = CurFrame->ChainedParent ? CurFrame->ChainedParent : CurFrame;
+
   CurFrame->ExceptionHandler = Sym;
   if (!Except && !Unwind)
     getContext().reportError(Loc, "Don't know what kind of handler this is!");
@@ -822,8 +824,6 @@ void MCStreamer::emitWinEHHandlerData(SMLoc Loc) {
   WinEH::FrameInfo *CurFrame = EnsureValidWinFrameInfo(Loc);
   if (!CurFrame)
     return;
-  if (CurFrame->ChainedParent)
-    getContext().reportError(Loc, "Chained unwind areas can't have handlers!");
 }
 
 void MCStreamer::emitCGProfileEntry(const MCSymbolRefExpr *From,
@@ -994,7 +994,8 @@ void MCStreamer::emitWinCFIBeginEpilogue(SMLoc Loc) {
   if (!CurFrame)
     return;
 
-  if (!CurFrame->PrologEnd)
+  // Chained unwinds aren't guaranteed to have a prolog.
+  if (!CurFrame->PrologEnd && !CurFrame->ChainedParent)
     return getContext().reportError(
         Loc, "starting epilogue (.seh_startepilogue) before prologue has ended "
              "(.seh_endprologue) in " +

llvm/lib/MC/MCWin64EH.cpp

@@ -22,8 +22,7 @@ class MCSection;
 
 /// MCExpr that represents the epilog unwind code in an unwind table.
 class MCUnwindV2EpilogTargetExpr final : public MCTargetExpr {
-  const MCSymbol *Function;
-  const MCSymbol *FunctionEnd;
+  const WinEH::FrameInfo &FrameInfo;
   const MCSymbol *UnwindV2Start;
   const MCSymbol *EpilogEnd;
   uint8_t EpilogSize;
@@ -32,9 +31,11 @@ class MCUnwindV2EpilogTargetExpr final : public MCTargetExpr {
   MCUnwindV2EpilogTargetExpr(const WinEH::FrameInfo &FrameInfo,
                              const WinEH::FrameInfo::Epilog &Epilog,
                              uint8_t EpilogSize_)
-      : Function(FrameInfo.Function), FunctionEnd(FrameInfo.FuncletOrFuncEnd),
-        UnwindV2Start(Epilog.UnwindV2Start), EpilogEnd(Epilog.End),
-        EpilogSize(EpilogSize_), Loc(Epilog.Loc) {}
+      : FrameInfo(FrameInfo), UnwindV2Start(Epilog.UnwindV2Start),
+        EpilogEnd(Epilog.End), EpilogSize(EpilogSize_), Loc(Epilog.Loc) {
+    assert(UnwindV2Start && "Epilog must have a start");
+    assert(EpilogEnd && "Epilog must have an end");
+  }
 
 public:
   static MCUnwindV2EpilogTargetExpr *
@@ -250,6 +251,10 @@ static void EmitUnwindInfo(MCStreamer &streamer, WinEH::FrameInfo *info) {
     // 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).
+    assert(
+        LastEpilog.UnwindV2Start &&
+        "If unwind v2 is enabled, epilog must have a unwind v2 start marker");
+    assert(LastEpilog.End && "Epilog must have an end");
     auto MaybeSize = GetOptionalAbsDifference(
         OS->getAssembler(), LastEpilog.End, LastEpilog.UnwindV2Start);
     if (!MaybeSize) {
@@ -270,9 +275,14 @@ static void EmitUnwindInfo(MCStreamer &streamer, WinEH::FrameInfo *info) {
     // 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);
+    // NOTE: At the point where the unwind info is emitted, the function may not
+    // have ended yet (e.g., if there is EH Handler Data), so assume that we
+    // aren't at the end (since we can't calculate it).
+    if (info->End) {
+      auto LastEpilogToFuncEnd = GetOptionalAbsDifference(
+          OS->getAssembler(), info->End, LastEpilog.UnwindV2Start);
+      LastEpilogIsAtEnd = (LastEpilogToFuncEnd == EpilogSize);
+    }
 
     // If we have an odd number of epilog codes, we need to add a padding code.
     size_t numEpilogCodes =
@@ -384,28 +394,28 @@ 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);
+  auto Offset = GetOptionalAbsDifference(*Asm, FrameInfo.End, UnwindV2Start);
   if (!Offset) {
     Asm->getContext().reportError(
         Loc, "Failed to evaluate epilog offset for Unwind v2 in " +
-                 Function->getName());
+                 FrameInfo.Function->getName());
     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 in " + Function->getName());
+        Loc, "Epilog offset is too large for Unwind v2 in " +
+                 FrameInfo.Function->getName());
     return false;
   }
 
-  // Sanity check that all epilogs are the same size.
+  // Validate 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->getName());
+                 FrameInfo.Function->getName());
     return false;
   }
 

llvm/lib/Target/X86/MCTargetDesc/X86WinCOFFStreamer.cpp

@@ -44,8 +44,11 @@ void X86WinCOFFStreamer::emitWinEHHandlerData(SMLoc Loc) {
 
   // We have to emit the unwind info now, because this directive
   // actually switches to the .xdata section.
-  if (WinEH::FrameInfo *CurFrame = getCurrentWinFrameInfo())
+  if (WinEH::FrameInfo *CurFrame = getCurrentWinFrameInfo()) {
+    // Handlers are always associated with the parent frame.
+    CurFrame = CurFrame->ChainedParent ? CurFrame->ChainedParent : CurFrame;
     EHStreamer.EmitUnwindInfo(*this, CurFrame, /* HandlerData = */ true);
+  }
 }
 
 void X86WinCOFFStreamer::emitWindowsUnwindTables(WinEH::FrameInfo *Frame) {

llvm/lib/Target/X86/X86InstrCompiler.td

@@ -272,6 +272,12 @@ let isPseudo = 1, isMeta = 1, SchedRW = [WriteSystem] in {
                             "#SEH_UnwindV2Start", []>;
 }
 
+// Chain instructions:
+let isPseudo = 1, isMeta = 1, SchedRW = [WriteSystem] in {
+  def SEH_SplitChained : I<0, Pseudo, (outs), (ins),
+                            "#SEH_SplitChained", []>;
+}
+
 //===----------------------------------------------------------------------===//
 // Pseudo instructions used by KCFI.
 //===----------------------------------------------------------------------===//

llvm/lib/Target/X86/X86MCInstLower.cpp

@@ -1795,6 +1795,10 @@ void X86AsmPrinter::EmitSEHInstruction(const MachineInstr *MI) {
     OutStreamer->emitWinCFIUnwindVersion(MI->getOperand(0).getImm());
     break;
 
+  case X86::SEH_SplitChained:
+    OutStreamer->emitWinCFISplitChained();
+    break;
+
   default:
     llvm_unreachable("expected SEH_ instruction");
   }
@@ -2526,6 +2530,7 @@ void X86AsmPrinter::emitInstruction(const MachineInstr *MI) {
   case X86::SEH_EndEpilogue:
   case X86::SEH_UnwindV2Start:
   case X86::SEH_UnwindVersion:
+  case X86::SEH_SplitChained:
     EmitSEHInstruction(MI);
     return;