Cranelift AArch64: Fix the get_return_address lowering (#4851)

The previous implementation assumed that nothing had clobbered the
LR register since the current function had started executing, so
it would be incorrect for a non-leaf function, for example, that
contains the `get_return_address` operation right after a call.
The operation is valid only if the `preserve_frame_pointers` flag
is enabled, which implies that the presence of a frame record on
the stack is guaranteed.

Copyright (c) 2022, Arm Limited.

cranelift/codegen/src/isa/aarch64/inst.isle

@@ -1714,6 +1714,9 @@
 (decl stack_reg () Reg)
 (extern constructor stack_reg stack_reg)
 
+(decl writable_link_reg () WritableReg)
+(extern constructor writable_link_reg writable_link_reg)
+
 (decl writable_zero_reg () WritableReg)
 (extern constructor writable_zero_reg writable_zero_reg)
 
@@ -2956,13 +2959,31 @@
 
 (decl aarch64_link () Reg)
 (rule (aarch64_link)
+      (if (preserve_frame_pointers))
       (if (sign_return_address_disabled))
-      (mov_preg (preg_link)))
+      (let ((dst WritableReg (temp_writable_reg $I64))
+            ;; Even though LR is not an allocatable register, whether it
+            ;; contains the return address for the current function is
+            ;; unknown at this point. For example, this operation may come
+            ;; immediately after a call, in which case LR would not have a
+            ;; valid value. That's why we must obtain the return address from
+            ;; the frame record that corresponds to the current subroutine on
+            ;; the stack; the presence of the record is guaranteed by the
+            ;; `preserve_frame_pointers` setting.
+            (addr AMode (AMode.FPOffset 8 $I64))
+            (_ Unit (emit (MInst.ULoad64 dst addr (mem_flags_trusted)))))
+           dst))
 
 (rule (aarch64_link)
-      ;; This constructor is always used for non-leaf functions, so it is safe
+      (if (preserve_frame_pointers))
-      ;; to clobber LR.
+      ;; Similarly to the rule above, we must load the return address from the
-      (let ((_ Unit (emit (MInst.Xpaclri))))
+      ;; the frame record. Furthermore, we can use LR as a scratch register
+      ;; because the function will set it to the return address immediately
+      ;; before returning.
+      (let ((addr AMode (AMode.FPOffset 8 $I64))
+            (lr WritableReg (writable_link_reg))
+            (_ Unit (emit (MInst.ULoad64 lr addr (mem_flags_trusted))))
+            (_ Unit (emit (MInst.Xpaclri))))
            (mov_preg (preg_link))))
 
 ;; Helper for getting the maximum shift amount for a type.

cranelift/codegen/src/isa/aarch64/lower/isle.rs

@@ -7,11 +7,12 @@ use generated_code::Context;
 // Types that the generated ISLE code uses via `use super::*`.
 use super::{
     fp_reg, lower_constant_f128, lower_constant_f32, lower_constant_f64, lower_fp_condcode,
-    stack_reg, writable_zero_reg, zero_reg, AMode, ASIMDFPModImm, ASIMDMovModImm, BranchTarget,
+    stack_reg, writable_link_reg, writable_zero_reg, zero_reg, AMode, ASIMDFPModImm,
-    CallIndInfo, CallInfo, Cond, CondBrKind, ExtendOp, FPUOpRI, FPUOpRIMod, FloatCC, Imm12,
+    ASIMDMovModImm, BranchTarget, CallIndInfo, CallInfo, Cond, CondBrKind, ExtendOp, FPUOpRI,
-    ImmLogic, ImmShift, Inst as MInst, IntCC, JTSequenceInfo, MachLabel, MemLabel, MoveWideConst,
+    FPUOpRIMod, FloatCC, Imm12, ImmLogic, ImmShift, Inst as MInst, IntCC, JTSequenceInfo,
-    MoveWideOp, NarrowValueMode, Opcode, OperandSize, PairAMode, Reg, SImm9, ScalarSize,
+    MachLabel, MemLabel, MoveWideConst, MoveWideOp, NarrowValueMode, Opcode, OperandSize,
-    ShiftOpAndAmt, UImm12Scaled, UImm5, VecMisc2, VectorSize, NZCV,
+    PairAMode, Reg, SImm9, ScalarSize, ShiftOpAndAmt, UImm12Scaled, UImm5, VecMisc2, VectorSize,
+    NZCV,
 };
 use crate::ir::condcodes;
 use crate::isa::aarch64::inst::{FPULeftShiftImm, FPURightShiftImm};
@@ -312,6 +313,10 @@ impl Context for IsleContext<'_, '_, MInst, Flags, IsaFlags, 6> {
         fp_reg()
     }
 
+    fn writable_link_reg(&mut self) -> WritableReg {
+        writable_link_reg()
+    }
+
     fn extended_value_from_value(&mut self, val: Value) -> Option<ExtendedValue> {
         let (val, extend) =
             super::get_as_extended_value(self.lower_ctx, val, NarrowValueMode::None)?;

cranelift/codegen/src/machinst/isle.rs

@@ -743,6 +743,15 @@ macro_rules! isle_prelude_methods {
             }
         }
 
+        #[inline]
+        fn preserve_frame_pointers(&mut self) -> Option<()> {
+            if self.flags.preserve_frame_pointers() {
+                Some(())
+            } else {
+                None
+            }
+        }
+
         #[inline]
         fn func_ref_data(&mut self, func_ref: FuncRef) -> (SigRef, ExternalName, RelocDistance) {
             let funcdata = &self.lower_ctx.dfg().ext_funcs[func_ref];

cranelift/codegen/src/prelude.isle

@@ -824,6 +824,9 @@
 (decl pure tls_model_is_coff () Unit)
 (extern constructor tls_model_is_coff tls_model_is_coff)
 
+(decl pure preserve_frame_pointers () Unit)
+(extern constructor preserve_frame_pointers preserve_frame_pointers)
+
 ;;;; Helpers for accessing instruction data ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 ;; Accessor for `FuncRef`.

cranelift/codegen/src/verifier/mod.rs

@@ -726,6 +726,8 @@ impl<'a> Verifier<'a> {
                 opcode: Opcode::GetFramePointer | Opcode::GetReturnAddress,
             } => {
                 if let Some(isa) = &self.isa {
+                    // Backends may already rely on this check implicitly, so do
+                    // not relax it without verifying that it is safe to do so.
                     if !isa.flags().preserve_frame_pointers() {
                         return errors.fatal((
                             inst,

cranelift/filetests/filetests/isa/aarch64/fp_sp_pc-pauth.clif

@@ -0,0 +1,47 @@
+test compile precise-output
+set preserve_frame_pointers=true
+target aarch64 sign_return_address
+
+function %fp() -> i64 {
+block0:
+    v0 = get_frame_pointer.i64
+    return v0
+}
+
+;   paciasp
+;   stp fp, lr, [sp, #-16]!
+;   mov fp, sp
+; block0:
+;   mov x0, fp
+;   ldp fp, lr, [sp], #16
+;   autiasp ; ret
+
+function %sp() -> i64 {
+block0:
+    v0 = get_stack_pointer.i64
+    return v0
+}
+
+;   paciasp
+;   stp fp, lr, [sp, #-16]!
+;   mov fp, sp
+; block0:
+;   mov x0, sp
+;   ldp fp, lr, [sp], #16
+;   autiasp ; ret
+
+function %return_address() -> i64 {
+block0:
+    v0 = get_return_address.i64
+    return v0
+}
+
+;   paciasp
+;   stp fp, lr, [sp, #-16]!
+;   mov fp, sp
+; block0:
+;   ldr lr, [fp, #8]
+;   xpaclri
+;   mov x0, lr
+;   ldp fp, lr, [sp], #16
+;   autiasp ; ret

cranelift/filetests/filetests/isa/aarch64/fp_sp_pc.clif

@@ -37,7 +37,6 @@ block0:
 ;   stp fp, lr, [sp, #-16]!
 ;   mov fp, sp
 ; block0:
-;   mov x0, lr
+;   ldr x0, [fp, #8]
 ;   ldp fp, lr, [sp], #16
 ;   ret
-