[MC] [Win64EH] Check that ARM64 prologs and epilogs have the right matching number of instructions

[mstorsjo]

Summary

This matches what was done for the ARM implementation (where getting the instruction sizes right is even more tricky, and hence needed tighter testing).

[mstorsjo]

Rebased, and fixed the wineh5.mir testcase which failed due to this patch.

With D134992, D135101 and D135103, I can manage to build large applications with this patch applied.

[mstorsjo]

I removed the RFC tag from this patch now, since all the other prerequisites have landed, and all of VLC builds successfully with this applied. So with this in place, we can notice other cases right away, where we generate prologues that don't match the unwind data.

[efriedma]

LGTM

llvm/lib/MC/MCWin64EH.cpp

 static int64_t GetAbsDifference(MCStreamer &Streamer, const MCSymbol *LHS,
                                 const MCSymbol *RHS) {
   Optional<int64_t> MaybeDiff = GetOptionalAbsDifference(Streamer, LHS, RHS);
   if (!MaybeDiff)
     report_fatal_error("Failed to evaluate function length in SEH unwind info");
   return *MaybeDiff;
 }
 
+static void checkARM64Instructions(MCStreamer &Streamer,
+                                   ArrayRef<WinEH::Instruction> Insns,
+                                   const MCSymbol *Begin, const MCSymbol *End,
+                                   StringRef Name, StringRef Type) {
+  if (!End)
+    return;
+  Optional<int64_t> MaybeDistance =
+      GetOptionalAbsDifference(Streamer, End, Begin);
+  if (!MaybeDistance)
+    return;
+  uint32_t Distance = (uint32_t)*MaybeDistance;
+
+  for (const auto &I : Insns) {
+    switch (static_cast<Win64EH::UnwindOpcodes>(I.Operation)) {
+    default:
+      break;
+    case Win64EH::UOP_TrapFrame:
+    case Win64EH::UOP_PushMachFrame:
+    case Win64EH::UOP_Context:
+    case Win64EH::UOP_ClearUnwoundToCall:
+      // Can't reason about these opcodes and how they map to actual
+      // instructions.
+      return;
+    }
+  }
+  // Exclude the end opcode which doesn't map to an instruction.
+  uint32_t InstructionBytes = 4 * (Insns.size() - 1);
+  if (Distance != InstructionBytes) {
+    Streamer.getContext().reportError(
+        SMLoc(), "Incorrect size for " + Name + " " + Type + ": " +
+                     Twine(Distance) +
+                     " bytes of instructions in range, but .seh directives "
+                     "corresponding to " +
+                     Twine(InstructionBytes) + " bytes\n");
+  }
+}
+
 static uint32_t ARM64CountOfUnwindCodes(ArrayRef<WinEH::Instruction> Insns) {
   uint32_t Count = 0;
   for (const auto &I : Insns) {
     switch (static_cast<Win64EH::UnwindOpcodes>(I.Operation)) {
     default:
       llvm_unreachable("Unsupported ARM64 unwind code");
     case Win64EH::UOP_AllocSmall:
       Count += 1;
if (MatchingEpilog) {
       AddedEpilogs.push_back(EpilogStart);
     }
   }
 }
 
 static void ARM64FindSegmentsInFunction(MCStreamer &streamer,
                                         WinEH::FrameInfo *info,
                                         int64_t RawFuncLength) {
+  if (info->PrologEnd)
+    checkARM64Instructions(streamer, info->Instructions, info->Begin,
+                           info->PrologEnd, info->Function->getName(),
+                           "prologue");
   struct EpilogStartEnd {
     MCSymbol *Start;
     int64_t Offset;
     int64_t End;
   };
   // Record Start and End of each epilog.
   SmallVector<struct EpilogStartEnd, 4> Epilogs;
   for (auto &I : info->EpilogMap) {
     MCSymbol *Start = I.first;
     auto &Instrs = I.second.Instructions;
     int64_t Offset = GetAbsDifference(streamer, Start, info->Begin);
+    checkARM64Instructions(streamer, Instrs, Start, I.second.End,
+                           info->Function->getName(), "epilogue");
     assert((Epilogs.size() == 0 || Offset >= Epilogs.back().End) &&
            "Epilogs should be monotonically ordered");
     // Exclue the end opcode from Instrs.size() when calculating the end of the
     // epilog.
     Epilogs.push_back({Start, Offset, Offset + (int64_t)(Instrs.size() - 1) * 4});
   }
 
   unsigned E = 0;

llvm/lib/Target/AArch64/MCTargetDesc/AArch64WinCOFFStreamer.cpp

void AArch64TargetWinCOFFStreamer::emitARM64WinCFIEpilogEnd() {
   WinEH::FrameInfo *CurFrame = S.EnsureValidWinFrameInfo(SMLoc());
   if (!CurFrame)
     return;
 
   InEpilogCFI = false;
   WinEH::Instruction Inst =
       WinEH::Instruction(Win64EH::UOP_End, /*Label=*/nullptr, -1, 0);
   CurFrame->EpilogMap[CurrentEpilog].Instructions.push_back(Inst);
+  MCSymbol *Label = S.emitCFILabel();
+  CurFrame->EpilogMap[CurrentEpilog].End = Label;
   CurrentEpilog = nullptr;
 }
 
 void AArch64TargetWinCOFFStreamer::emitARM64WinCFITrapFrame() {
   emitARM64WinUnwindCode(Win64EH::UOP_TrapFrame, -1, 0);
 }
 
 void AArch64TargetWinCOFFStreamer::emitARM64WinCFIMachineFrame() {

llvm/test/CodeGen/AArch64/wineh5.mir

 # RUN: llc -o - %s -mtriple=aarch64-windows -start-after=prologepilog \
 # RUN:    -filetype=obj | llvm-readobj --unwind - | FileCheck %s
 
 # Check that that the large stack allocation is correctly represented in .xdata.
 
 # CHECK:        ExceptionData {
 # CHECK-NEXT:     FunctionLength: 156
 # CHECK-NEXT:     Version: 0
 # CHECK-NEXT:     ExceptionData: No
 # CHECK-NEXT:     EpiloguePacked: Yes
-# CHECK-NEXT:     EpilogueOffset: 10
-# CHECK-NEXT:     ByteCodeLength: 20
+# CHECK-NEXT:     EpilogueOffset: 11
+# CHECK-NEXT:     ByteCodeLength: 24
 # CHECK-NEXT:     Prologue [
 # CHECK-NEXT:       0xe002dac9          ; sub sp, #2993296
 # CHECK-NEXT:       0xe3                ; nop
 # CHECK-NEXT:       0xe3                ; nop
+# CHECK-NEXT:       0xe3                ; nop
 # CHECK-NEXT:       0x42                ; stp x29, x30, [sp, #16]
 # CHECK-NEXT:       0xd53f              ; str x28, [sp, #-256]!
 # CHECK-NEXT:       0xe4                ; end
 # CHECK-NEXT:     ]
 # CHECK-NEXT:     Epilogue [
 # CHECK-NEXT:       0xe002da00          ; add sp, #2990080
 # CHECK-NEXT:       0xc0c9              ; add sp, #3216
 # CHECK-NEXT:       0x42                ; ldp x29, x30, [sp, #16]
body:             |
   bb.1.entry:
     successors: %bb.2(0x40000000), %bb.3(0x40000000)
     liveins: $w0, $x28, $fp, $lr
 
     early-clobber $sp = frame-setup STRXpre killed $x28, $sp, -32 :: (store (s64) into %stack.9)
     frame-setup SEH_SaveReg_X 28, -256
     frame-setup STPXi killed $fp, killed $lr, $sp, 2 :: (store (s64) into %stack.7), (store (s64) into %stack.8)
     frame-setup SEH_SaveFPLR 16
-    $x15 = frame-setup MOVi64imm 187081
+    $x15 = frame-setup MOVZXi 56009, 0
+    frame-setup SEH_Nop
+    $x15 = frame-setup MOVKXi $x15, 2, 16
     frame-setup SEH_Nop
     frame-setup BL &__chkstk, implicit-def $lr, implicit $sp, implicit $x15
     frame-setup SEH_Nop
     $sp = frame-setup SUBXrx64 killed $sp, killed $x15, 28
     frame-setup SEH_StackAlloc 2993296
     frame-setup SEH_PrologEnd
     $x8 = ADDXri $sp, 730, 12
     $x8 = ADDXri $x8, 3208, 0