[MC] [Win64EH] Generate ARM64 packed unwind info with signed return addresses

[efriedma]

LGTM

llvm/lib/MC/MCWin64EH.cpp

if (PackedEpilogOffset == 0) {
     // Too much difference between prolog and epilog.
     return false;
   }
   unsigned RegI = 0, RegF = 0;
   int Predecrement = 0;
   enum {
     Start,
     Start2,
+    Start3,
     IntRegs,
     FloatRegs,
     InputArgs,
     StackAdjust,
     FrameRecord,
     End
   } Location = Start;
   bool StandaloneLR = false, FPLRPair = false;
+  bool PAC = false;
   int StackOffset = 0;
   int Nops = 0;
   // Iterate over the prolog and check that all opcodes exactly match
   // the canonical order and form. A more lax check could verify that
   // all saved registers are in the expected locations, but not enforce
   // the order - that would work fine when unwinding from within
   // functions, but not be exactly right if unwinding happens within
   // prologs/epilogs.
   for (const WinEH::Instruction &Inst : info->Instructions) {
     switch (Inst.Operation) {
     case Win64EH::UOP_End:
       if (Location != Start)
         return false;
       Location = Start2;
       break;
-    case Win64EH::UOP_SaveR19R20X:
+    case Win64EH::UOP_PACSignReturnAddress:
       if (Location != Start2)
         return false;
+      PAC = true;
+      Location = Start3;
+      break;
+    case Win64EH::UOP_SaveR19R20X:
+      if (Location != Start2 && Location != Start3)
+        return false;
       Predecrement = Inst.Offset;
       RegI = 2;
       Location = IntRegs;
       break;
     case Win64EH::UOP_SaveRegX:
-      if (Location != Start2)
+      if (Location != Start2 && Location != Start3)
         return false;
       Predecrement = Inst.Offset;
       if (Inst.Register == 19)
         RegI += 1;
       else if (Inst.Register == 30)
         StandaloneLR = true;
       else
         return false;
for (const WinEH::Instruction &Inst : info->Instructions) {
     case Win64EH::UOP_SaveFReg:
       if (Location != FloatRegs || RegF == 0 || Inst.Register != 8 + RegF ||
           Inst.Offset != 8 * (RegI + (StandaloneLR ? 1 : 0) + RegF))
         return false;
       RegF += 1;
       Location = InputArgs;
       break;
     case Win64EH::UOP_SaveFRegPX:
-      if (Location != Start2 || Inst.Register != 8)
+      if ((Location != Start2 && Location != Start3) || Inst.Register != 8)
         return false;
       Predecrement = Inst.Offset;
       RegF = 2;
       Location = FloatRegs;
       break;
     case Win64EH::UOP_SaveFRegP:
       if ((Location != IntRegs && Location != FloatRegs) ||
           Inst.Register != 8 + RegF ||
for (const WinEH::Instruction &Inst : info->Instructions) {
     case Win64EH::UOP_Nop:
       if (Location != IntRegs && Location != FloatRegs && Location != InputArgs)
         return false;
       Location = InputArgs;
       Nops++;
       break;
     case Win64EH::UOP_AllocSmall:
     case Win64EH::UOP_AllocMedium:
-      if (Location != Start2 && Location != IntRegs && Location != FloatRegs &&
-          Location != InputArgs && Location != StackAdjust)
+      if (Location != Start2 && Location != Start3 && Location != IntRegs &&
+          Location != FloatRegs && Location != InputArgs &&
+          Location != StackAdjust)
         return false;
       // Can have either a single decrement, or a pair of decrements with
       // 4080 and another decrement.
       if (StackOffset == 0)
         StackOffset = Inst.Offset;
       else if (StackOffset != 4080)
         return false;
       else
         StackOffset += Inst.Offset;
       Location = StackAdjust;
       break;
     case Win64EH::UOP_SaveFPLRX:
       // Not allowing FPLRX after StackAdjust; if a StackAdjust is used, it
       // should be followed by a FPLR instead.
-      if (Location != Start2 && Location != IntRegs && Location != FloatRegs &&
-          Location != InputArgs)
+      if (Location != Start2 && Location != Start3 && Location != IntRegs &&
+          Location != FloatRegs && Location != InputArgs)
         return false;
       StackOffset = Inst.Offset;
       Location = FrameRecord;
       FPLRPair = true;
       break;
     case Win64EH::UOP_SaveFPLR:
       // This can only follow after a StackAdjust
       if (Location != StackAdjust || Inst.Offset != 0)
static bool tryARM64PackedUnwind(WinEH::FrameInfo *info, uint32_t FuncLength,
   if (RegI > 10 || RegF > 8)
     return false;
   if (StandaloneLR && FPLRPair)
     return false;
   if (FPLRPair && Location != End)
     return false;
   if (Nops != 0 && Nops != 4)
     return false;
+  if (PAC && !FPLRPair)
+    return false;
   int H = Nops == 4;
   // There's an inconsistency regarding packed unwind info with homed
   // parameters; according to the documentation, the epilog shouldn't have
   // the same corresponding nops (and thus, to set the H bit, we should
   // require an epilog which isn't exactly symmetrical - we shouldn't accept
   // an exact mirrored epilog for those cases), but in practice,
   // RtlVirtualUnwind behaves as if it does expect the epilogue to contain
   // the same nops. See https://github.com/llvm/llvm-project/issues/54879.
static bool tryARM64PackedUnwind(WinEH::FrameInfo *info, uint32_t FuncLength,
   uint32_t FrameSize = (StackOffset + SavSZ) / 16;
   if (FrameSize > 0x1FF)
     return false;
   assert(RegF != 1 && "One single float reg not allowed");
   if (RegF > 0)
     RegF--; // Convert from actual number of registers, to value stored
   assert(FuncLength <= 0x7FF && "FuncLength should have been checked earlier");
   int Flag = 0x01; // Function segments not supported yet
-  int CR = FPLRPair ? 3 : StandaloneLR ? 1 : 0;
+  int CR = PAC ? 2 : FPLRPair ? 3 : StandaloneLR ? 1 : 0;
   info->PackedInfo |= Flag << 0;
   info->PackedInfo |= (FuncLength & 0x7FF) << 2;
   info->PackedInfo |= (RegF & 0x7) << 13;
   info->PackedInfo |= (RegI & 0xF) << 16;
   info->PackedInfo |= (H & 0x1) << 20;
   info->PackedInfo |= (CR & 0x3) << 21;
   info->PackedInfo |= (FrameSize & 0x1FF) << 23;
   return true;

llvm/test/MC/AArch64/seh-packed-unwind.s

 // CHECK-NEXT:     CR: 3
 // CHECK-NEXT:     FrameSize: 32
 // CHECK-NEXT:     Prologue [
 // CHECK-NEXT:       mov x29, sp
 // CHECK-NEXT:       stp x29, lr, [sp, #-32]!
 // CHECK-NEXT:       end
 // CHECK-NEXT:     ]
 // CHECK-NEXT:   }
+// CHECK-NEXT:   RuntimeFunction {
+// CHECK-NEXT:     Function: func16
+// CHECK-NEXT:     Fragment: No
+// CHECK-NEXT:     FunctionLength: 28
+// CHECK-NEXT:     RegF: 0
+// CHECK-NEXT:     RegI: 0
+// CHECK-NEXT:     HomedParameters: No
+// CHECK-NEXT:     CR: 2
+// CHECK-NEXT:     FrameSize: 32
+// CHECK-NEXT:     Prologue [
+// CHECK-NEXT:       mov x29, sp
+// CHECK-NEXT:       stp x29, lr, [sp, #-32]!
+// CHECK-NEXT:       pacibsp
+// CHECK-NEXT:       end
+// CHECK-NEXT:     ]
+// CHECK-NEXT:   }
+// CHECK-NEXT:   RuntimeFunction {
+// CHECK-NEXT:     Function: func17
+// CHECK-NEXT:     Fragment: No
+// CHECK-NEXT:     FunctionLength: 40
+// CHECK-NEXT:     RegF: 0
+// CHECK-NEXT:     RegI: 2
+// CHECK-NEXT:     HomedParameters: No
+// CHECK-NEXT:     CR: 2
+// CHECK-NEXT:     FrameSize: 48
+// CHECK-NEXT:     Prologue [
+// CHECK-NEXT:       mov x29, sp
+// CHECK-NEXT:       stp x29, lr, [sp, #-32]!
+// CHECK-NEXT:       stp x19, x20, [sp, #-16]!
+// CHECK-NEXT:       pacibsp
+// CHECK-NEXT:       end
+// CHECK-NEXT:     ]
+// CHECK-NEXT:   }
+// CHECK-NEXT:   RuntimeFunction {
+// CHECK-NEXT:     Function: func18
+// CHECK-NEXT:     Fragment: No
+// CHECK-NEXT:     FunctionLength: 56
+// CHECK-NEXT:     RegF: 0
+// CHECK-NEXT:     RegI: 2
+// CHECK-NEXT:     HomedParameters: No
+// CHECK-NEXT:     CR: 2
+// CHECK-NEXT:     FrameSize: 4112
+// CHECK-NEXT:     Prologue [
+// CHECK-NEXT:       mov x29, sp
+// CHECK-NEXT:       stp x29, lr, [sp, #0]
+// CHECK-NEXT:       sub sp, sp, #16
+// CHECK-NEXT:       sub sp, sp, #4080
+// CHECK-NEXT:       stp x19, x20, [sp, #-16]!
+// CHECK-NEXT:       pacibsp
+// CHECK-NEXT:       end
+// CHECK-NEXT:     ]
+// CHECK-NEXT:   }
 // CHECK:        RuntimeFunction {
 // CHECK-NEXT:     Function: nonpacked1
 // CHECK-NEXT:     ExceptionRecord:
 // CHECK-NEXT:     ExceptionData {
 // CHECK:            EpiloguePacked: Yes
 // CHECK:        RuntimeFunction {
 // CHECK-NEXT:     Function: nonpacked2
 // CHECK-NEXT:     ExceptionRecord:

 // CHECK-NEXT:     ExceptionRecord:
 // CHECK-NEXT:     ExceptionData {
 // CHECK:            EpiloguePacked: Yes
 // CHECK:        RuntimeFunction {
 // CHECK-NEXT:     Function: nonpacked13
 // CHECK-NEXT:     ExceptionRecord:
 // CHECK-NEXT:     ExceptionData {
 // CHECK:            EpiloguePacked: Yes
+// CHECK:        RuntimeFunction {
+// CHECK-NEXT:     Function: nonpacked14
+// CHECK-NEXT:     ExceptionRecord:
+// CHECK-NEXT:     ExceptionData {
+// CHECK:            EpiloguePacked: Yes
+// CHECK:        RuntimeFunction {
+// CHECK-NEXT:     Function: nonpacked15
+// CHECK-NEXT:     ExceptionRecord:
+// CHECK-NEXT:     ExceptionData {
+// CHECK:            EpiloguePacked: Yes
 
     .text
 func1:
     .seh_proc func1
     stp x19, x20, [sp, #-144]!
     .seh_save_r19r20_x 144
     stp x21, x22, [sp, #16]
     .seh_save_regp x21, 16
func15:
     .seh_startepilogue
     // Epilogue missing the .seh_set_fp, but still generating packed info.
     ldp x29, lr,  [sp], #32
     .seh_save_fplr_x 32
     .seh_endepilogue
     ret
     .seh_endproc
 
+func16:
+    .seh_proc func16
+    pacibsp
+    .seh_pac_sign_return_address
+    stp x29, lr,  [sp, #-32]!
+    .seh_save_fplr_x 32
+    mov x29, sp
+    .seh_set_fp
+    .seh_endprologue
+    nop
+    .seh_startepilogue
+    ldp x29, lr, [sp], #32
+    .seh_save_fplr_x 32
+    autibsp
+    .seh_pac_sign_return_address
+    .seh_endepilogue
+    ret
+    .seh_endproc
+
+func17:
+    .seh_proc func17
+    pacibsp
+    .seh_pac_sign_return_address
+    stp x19, x20, [sp, #-16]!
+    .seh_save_r19r20_x 16
+    stp x29, lr,  [sp, #-32]!
+    .seh_save_fplr_x 32
+    mov x29, sp
+    .seh_set_fp
+    .seh_endprologue
+    nop
+    .seh_startepilogue
+    mov sp,  x29
+    .seh_set_fp
+    ldp x29, lr,  [sp], #32
+    .seh_save_fplr_x 32
+    ldp x19, x20, [sp], #16
+    .seh_save_r19r20_x 16
+    autibsp
+    .seh_pac_sign_return_address
+    .seh_endepilogue
+    ret
+    .seh_endproc
+
+func18:
+    .seh_proc func18
+    pacibsp
+    .seh_pac_sign_return_address
+    stp x19, x20, [sp, #-16]!
+    .seh_save_r19r20_x 16
+    sub sp,  sp,  #4080
+    .seh_stackalloc 4080
+    sub sp,  sp,  #16
+    .seh_stackalloc 16
+    stp x29, lr,  [sp, #0]
+    .seh_save_fplr 0
+    mov x29, sp
+    .seh_set_fp
+    .seh_endprologue
+    nop
+    .seh_startepilogue
+    mov sp,  x29
+    .seh_set_fp
+    ldp x29, lr,  [sp, #0]
+    .seh_save_fplr 0
+    add sp,  sp,  #16
+    .seh_stackalloc 16
+    add sp,  sp,  #4080
+    .seh_stackalloc 4080
+    ldp x19, x20, [sp], #16
+    .seh_save_r19r20_x 16
+    autibsp
+    .seh_pac_sign_return_address
+    .seh_endepilogue
+    ret
+    .seh_endproc
+
 nonpacked1:
     .seh_proc nonpacked1
     // Can't be packed; can't save integer registers after float registers.
     stp d8,  d9,  [sp, #-32]!
     .seh_save_fregp_x d8, 32
     stp x19, x20, [sp, #16]!
     .seh_save_regp x19, 16
     .seh_endprologue
nonpacked13:
     .seh_stackalloc 16
     .seh_endepilogue
     ret
     .seh_endfunclet
     .seh_handlerdata
     .long 0
     .text
     .seh_endproc
+
+nonpacked14:
+    .seh_proc nonpacked14
+    // Can't be packed; a signed return address can only be expressed if
+    // we save both x29 and lr on the stack.
+    pacibsp
+    .seh_pac_sign_return_address
+    str lr,       [sp, #-32]!
+    .seh_save_reg_x lr, 32
+    .seh_endprologue
+    nop
+    .seh_startepilogue
+    ldr lr,       [sp], #32
+    .seh_save_reg_x lr, 32
+    autibsp
+    .seh_pac_sign_return_address
+    .seh_endepilogue
+    ret
+    .seh_endproc
+
+nonpacked15:
+    .seh_proc nonpacked15
+    // Can't be packed; a signed return address can only be expressed if
+    // we save both x29 and lr on the stack.
+    pacibsp
+    .seh_pac_sign_return_address
+    stp x19, x20, [sp, #-32]!
+    .seh_save_r19r20_x 32
+    stp x21, lr,  [sp, #16]
+    .seh_save_lrpair x21, 16
+    sub sp,  sp,  #16
+    .seh_stackalloc 16
+    .seh_endprologue
+    nop
+    .seh_startepilogue
+    add sp,  sp,  #16
+    .seh_stackalloc 16
+    ldp x21, lr,  [sp, #16]
+    .seh_save_lrpair x21, 16
+    ldp x19, x20, [sp], #32
+    .seh_save_r19r20_x 32
+    autibsp
+    .seh_pac_sign_return_address
+    .seh_endepilogue
+    ret
+    .seh_endproc