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

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

This will allow catching any future cases where prologs and epilogs
don't match the instructions within them.

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

Author: mstorsjo

llvm/lib/MC/MCWin64EH.cpp

@@ -299,6 +299,43 @@ static int64_t GetAbsDifference(MCStreamer &Streamer, const MCSymbol *LHS,
   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) {
@@ -1002,6 +1039,10 @@ static void ARM64ProcessEpilogs(WinEH::FrameInfo *info,
 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;
@@ -1013,6 +1054,8 @@ static void ARM64FindSegmentsInFunction(MCStreamer &streamer,
     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

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

@@ -202,6 +202,8 @@ void AArch64TargetWinCOFFStreamer::emitARM64WinCFIEpilogEnd() {
   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;
 }
 

llvm/test/CodeGen/AArch64/wineh5.mir

@@ -8,12 +8,13 @@
 # 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
@@ -158,7 +159,9 @@ body:             |
     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