rustc_codegen_llvm: Add a new 'pc' option to branch-protection

Add a new 'pc' option to -Z branch-protection for aarch64 that
enables the use of PC as a diversifier in PAC branch protection code.

When the pauth-lr target feature is enabled in combination
with -Z branch-protection=pac-ret,pc, the new 9.5-a instructions
(pacibsppc, retaasppc, etc) will be generated.

12 files changed, 97 insertions, 17 deletions

diff --git a/compiler/rustc_codegen_llvm/src/attributes.rs b/compiler/rustc_codegen_llvm/src/attributes.rs
index 2c5ec9dad59..2af90142d2f 100644
--- a/compiler/rustc_codegen_llvm/src/attributes.rs
+++ b/compiler/rustc_codegen_llvm/src/attributes.rs
@@ -424,7 +424,10 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
                 if bti {
                     to_add.push(llvm::CreateAttrString(cx.llcx, "branch-target-enforcement"));
                 }
-                if let Some(PacRet { leaf, key }) = pac_ret {
+                if let Some(PacRet { leaf, pc, key }) = pac_ret {
+                    if pc {
+                        to_add.push(llvm::CreateAttrString(cx.llcx, "branch-protection-pauth-lr"));
+                    }
                     to_add.push(llvm::CreateAttrStringValue(
                         cx.llcx,
                         "sign-return-address",
diff --git a/compiler/rustc_codegen_llvm/src/context.rs b/compiler/rustc_codegen_llvm/src/context.rs
index 9778ff4918c..924137e02b1 100644
--- a/compiler/rustc_codegen_llvm/src/context.rs
+++ b/compiler/rustc_codegen_llvm/src/context.rs
@@ -302,7 +302,13 @@ pub(crate) unsafe fn create_module<'ll>(
                 "sign-return-address",
                 pac_ret.is_some().into(),
             );
-            let pac_opts = pac_ret.unwrap_or(PacRet { leaf: false, key: PAuthKey::A });
+            let pac_opts = pac_ret.unwrap_or(PacRet { leaf: false, pc: false, key: PAuthKey::A });
+            llvm::add_module_flag_u32(
+                llmod,
+                llvm::ModuleFlagMergeBehavior::Min,
+                "branch-protection-pauth-lr",
+                pac_opts.pc.into(),
+            );
             llvm::add_module_flag_u32(
                 llmod,
                 llvm::ModuleFlagMergeBehavior::Min,
diff --git a/compiler/rustc_interface/src/tests.rs b/compiler/rustc_interface/src/tests.rs
index d3762e739db..546addf2232 100644
--- a/compiler/rustc_interface/src/tests.rs
+++ b/compiler/rustc_interface/src/tests.rs
@@ -764,7 +764,7 @@ fn test_unstable_options_tracking_hash() {
         branch_protection,
         Some(BranchProtection {
             bti: true,
-            pac_ret: Some(PacRet { leaf: true, key: PAuthKey::B })
+            pac_ret: Some(PacRet { leaf: true, pc: true, key: PAuthKey::B })
         })
     );
     tracked!(codegen_backend, Some("abc".to_string()));
diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs
index d733e32f209..66312d300da 100644
--- a/compiler/rustc_session/src/config.rs
+++ b/compiler/rustc_session/src/config.rs
@@ -1319,6 +1319,7 @@ pub enum PAuthKey {
 #[derive(Clone, Copy, Hash, Debug, PartialEq)]
 pub struct PacRet {
     pub leaf: bool,
+    pub pc: bool,
     pub key: PAuthKey,
 }
 
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
index 54a4621db24..8eba4c68990 100644
--- a/compiler/rustc_session/src/options.rs
+++ b/compiler/rustc_session/src/options.rs
@@ -442,8 +442,7 @@ mod desc {
     pub(crate) const parse_polonius: &str = "either no value or `legacy` (the default), or `next`";
     pub(crate) const parse_stack_protector: &str =
         "one of (`none` (default), `basic`, `strong`, or `all`)";
-    pub(crate) const parse_branch_protection: &str =
-        "a `,` separated combination of `bti`, `b-key`, `pac-ret`, or `leaf`";
+    pub(crate) const parse_branch_protection: &str = "a `,` separated combination of `bti`, `pac-ret`, followed by a combination of `pc`, `b-key`, or `leaf`";
     pub(crate) const parse_proc_macro_execution_strategy: &str =
         "one of supported execution strategies (`same-thread`, or `cross-thread`)";
     pub(crate) const parse_remap_path_scope: &str =
@@ -1396,7 +1395,7 @@ mod parse {
                     match opt {
                         "bti" => slot.bti = true,
                         "pac-ret" if slot.pac_ret.is_none() => {
-                            slot.pac_ret = Some(PacRet { leaf: false, key: PAuthKey::A })
+                            slot.pac_ret = Some(PacRet { leaf: false, pc: false, key: PAuthKey::A })
                         }
                         "leaf" => match slot.pac_ret.as_mut() {
                             Some(pac) => pac.leaf = true,
@@ -1406,6 +1405,10 @@ mod parse {
                             Some(pac) => pac.key = PAuthKey::B,
                             _ => return false,
                         },
+                        "pc" => match slot.pac_ret.as_mut() {
+                            Some(pac) => pac.pc = true,
+                            _ => return false,
+                        },
                         _ => return false,
                     };
                 }
diff --git a/src/doc/unstable-book/src/compiler-flags/branch-protection.md b/src/doc/unstable-book/src/compiler-flags/branch-protection.md
index 9276220f447..f0cc44a07f3 100644
--- a/src/doc/unstable-book/src/compiler-flags/branch-protection.md
+++ b/src/doc/unstable-book/src/compiler-flags/branch-protection.md
@@ -9,11 +9,12 @@ This option is only accepted when targeting AArch64 architectures.
 It takes some combination of the following values, separated by a `,`.
 
 - `pac-ret` - Enable pointer authentication for non-leaf functions.
+- `pc` - Use PC as a diversifier using PAuthLR instructions
 - `leaf` - Enable pointer authentication for all functions, including leaf functions.
 - `b-key` - Sign return addresses with key B, instead of the default key A.
 - `bti` - Enable branch target identification.
 
-`leaf` and `b-key` are only valid if `pac-ret` was previously specified.
+`leaf`, `b-key` and `pc` are only valid if `pac-ret` was previously specified.
 For example, `-Z branch-protection=bti,pac-ret,leaf` is valid, but
 `-Z branch-protection=bti,leaf,pac-ret` is not.
 
diff --git a/tests/assembly/aarch64-pointer-auth.rs b/tests/assembly/aarch64-pointer-auth.rs
index 1e53878a2cc..344e9e74bc2 100644
--- a/tests/assembly/aarch64-pointer-auth.rs
+++ b/tests/assembly/aarch64-pointer-auth.rs
@@ -1,9 +1,13 @@
 // Test that PAC instructions are emitted when branch-protection is specified.
 
+//@ revisions: PACRET PAUTHLR_NOP PAUTHLR
 //@ assembly-output: emit-asm
-//@ compile-flags: --target aarch64-unknown-linux-gnu
-//@ compile-flags: -Z branch-protection=pac-ret,leaf
 //@ needs-llvm-components: aarch64
+//@ compile-flags: --target aarch64-unknown-linux-gnu
+//@ [PACRET] compile-flags: -Z branch-protection=pac-ret,leaf
+//@ [PAUTHLR_NOP] compile-flags: -Z branch-protection=pac-ret,pc,leaf
+//@ [PAUTHLR] compile-flags: -C target-feature=+pauth-lr -Z branch-protection=pac-ret,pc,leaf
+//@ min-llvm-version: 19
 
 #![feature(no_core, lang_items)]
 #![no_std]
@@ -13,8 +17,13 @@
 #[lang = "sized"]
 trait Sized {}
 
-// CHECK: hint #25
-// CHECK: hint #29
+// PACRET: hint #25
+// PACRET: hint #29
+// PAUTHLR_NOP: hint #25
+// PAUTHLR_NOP: hint #39
+// PAUTHLR_NOP: hint #29
+// PAUTHLR: paciasppc
+// PAUTHLR: autiasppc
 #[no_mangle]
 pub fn test() -> u8 {
     42
diff --git a/tests/codegen/branch-protection.rs b/tests/codegen/branch-protection.rs
index 2f5ff9e98c2..945bad05625 100644
--- a/tests/codegen/branch-protection.rs
+++ b/tests/codegen/branch-protection.rs
@@ -1,11 +1,15 @@
 // Test that the correct module flags are emitted with different branch protection flags.
 
-//@ revisions: BTI PACRET LEAF BKEY NONE
+//@ revisions: BTI PACRET LEAF BKEY PAUTHLR PAUTHLR_BKEY PAUTHLR_LEAF PAUTHLR_BTI NONE
 //@ needs-llvm-components: aarch64
 //@ [BTI] compile-flags: -Z branch-protection=bti
 //@ [PACRET] compile-flags: -Z branch-protection=pac-ret
 //@ [LEAF] compile-flags: -Z branch-protection=pac-ret,leaf
 //@ [BKEY] compile-flags: -Z branch-protection=pac-ret,b-key
+//@ [PAUTHLR] compile-flags: -Z branch-protection=pac-ret,pc
+//@ [PAUTHLR_BKEY] compile-flags: -Z branch-protection=pac-ret,pc,b-key
+//@ [PAUTHLR_LEAF] compile-flags: -Z branch-protection=pac-ret,pc,leaf
+//@ [PAUTHLR_BTI] compile-flags: -Z branch-protection=bti,pac-ret,pc
 //@ compile-flags: --target aarch64-unknown-linux-gnu
 //@ min-llvm-version: 19
 
@@ -24,6 +28,7 @@ pub fn test() {}
 // BTI: attributes [[ATTR]] = {{.*}} "branch-target-enforcement"
 // BTI: !"branch-target-enforcement", i32 1
 // BTI: !"sign-return-address", i32 0
+// BTI: !"branch-protection-pauth-lr", i32 0
 // BTI: !"sign-return-address-all", i32 0
 // BTI: !"sign-return-address-with-bkey", i32 0
 
@@ -31,6 +36,7 @@ pub fn test() {}
 // PACRET-SAME: "sign-return-address-key"="a_key"
 // PACRET: !"branch-target-enforcement", i32 0
 // PACRET: !"sign-return-address", i32 1
+// PACRET: !"branch-protection-pauth-lr", i32 0
 // PACRET: !"sign-return-address-all", i32 0
 // PACRET: !"sign-return-address-with-bkey", i32 0
 
@@ -38,6 +44,7 @@ pub fn test() {}
 // LEAF-SAME: "sign-return-address-key"="a_key"
 // LEAF: !"branch-target-enforcement", i32 0
 // LEAF: !"sign-return-address", i32 1
+// LEAF: !"branch-protection-pauth-lr", i32 0
 // LEAF: !"sign-return-address-all", i32 1
 // LEAF: !"sign-return-address-with-bkey", i32 0
 
@@ -45,9 +52,42 @@ pub fn test() {}
 // BKEY-SAME: "sign-return-address-key"="b_key"
 // BKEY: !"branch-target-enforcement", i32 0
 // BKEY: !"sign-return-address", i32 1
+// BKEY: !"branch-protection-pauth-lr", i32 0
 // BKEY: !"sign-return-address-all", i32 0
 // BKEY: !"sign-return-address-with-bkey", i32 1
 
+// PAUTHLR: attributes [[ATTR]] = {{.*}} "sign-return-address"="non-leaf"
+// PAUTHLR-SAME: "sign-return-address-key"="a_key"
+// PAUTHLR: !"branch-target-enforcement", i32 0
+// PAUTHLR: !"sign-return-address", i32 1
+// PAUTHLR: !"branch-protection-pauth-lr", i32 1
+// PAUTHLR: !"sign-return-address-all", i32 0
+// PAUTHLR: !"sign-return-address-with-bkey", i32 0
+
+// PAUTHLR_BKEY: attributes [[ATTR]] = {{.*}} "sign-return-address"="non-leaf"
+// PAUTHLR_BKEY-SAME: "sign-return-address-key"="b_key"
+// PAUTHLR_BKEY: !"branch-target-enforcement", i32 0
+// PAUTHLR_BKEY: !"sign-return-address", i32 1
+// PAUTHLR_BKEY: !"branch-protection-pauth-lr", i32 1
+// PAUTHLR_BKEY: !"sign-return-address-all", i32 0
+// PAUTHLR_BKEY: !"sign-return-address-with-bkey", i32 1
+
+// PAUTHLR_LEAF: attributes [[ATTR]] = {{.*}} "sign-return-address"="all"
+// PAUTHLR_LEAF-SAME: "sign-return-address-key"="a_key"
+// PAUTHLR_LEAF: !"branch-target-enforcement", i32 0
+// PAUTHLR_LEAF: !"sign-return-address", i32 1
+// PAUTHLR_LEAF: !"branch-protection-pauth-lr", i32 1
+// PAUTHLR_LEAF: !"sign-return-address-all", i32 1
+// PAUTHLR_LEAF: !"sign-return-address-with-bkey", i32 0
+
+// PAUTHLR_BTI: attributes [[ATTR]] = {{.*}} "sign-return-address"="non-leaf"
+// PAUTHLR_BTI-SAME: "sign-return-address-key"="a_key"
+// PAUTHLR_BTI: !"branch-target-enforcement", i32 1
+// PAUTHLR_BTI: !"sign-return-address", i32 1
+// PAUTHLR_BTI: !"branch-protection-pauth-lr", i32 1
+// PAUTHLR_BTI: !"sign-return-address-all", i32 0
+// PAUTHLR_BTI: !"sign-return-address-with-bkey", i32 0
+
 // NONE-NOT: branch-target-enforcement
 // NONE-NOT: sign-return-address
 // NONE-NOT: sign-return-address-all
diff --git a/tests/run-make/pointer-auth-link-with-c/rmake.rs b/tests/run-make/pointer-auth-link-with-c/rmake.rs
index 960eafa546b..7b6dff10eae 100644
--- a/tests/run-make/pointer-auth-link-with-c/rmake.rs
+++ b/tests/run-make/pointer-auth-link-with-c/rmake.rs
@@ -1,7 +1,7 @@
 // `-Z branch protection` is an unstable compiler feature which adds pointer-authentication
 // code (PAC), a useful hashing measure for verifying that pointers have not been modified.
 // This test checks that compilation and execution is successful when this feature is activated,
-// with some of its possible extra arguments (bti, pac-ret, leaf).
+// with some of its possible extra arguments (bti, pac-ret, pc, leaf, b-key).
 // See https://github.com/rust-lang/rust/pull/88354
 
 //@ only-aarch64
@@ -25,4 +25,16 @@ fn main() {
     llvm_ar().obj_to_ar().output_input("libtest.a", &obj_file).run();
     rustc().arg("-Zbranch-protection=bti,pac-ret,leaf").input("test.rs").run();
     run("test");
+
+    // FIXME: +pc was only recently added to LLVM
+    // cc().arg("-v")
+    //     .arg("-c")
+    //     .out_exe("test")
+    //     .input("test.c")
+    //     .arg("-mbranch-protection=bti+pac-ret+pc+leaf")
+    //     .run();
+    // let obj_file = if is_msvc() { "test.obj" } else { "test" };
+    // llvm_ar().obj_to_ar().output_input("libtest.a", &obj_file).run();
+    // rustc().arg("-Zbranch-protection=bti,pac-ret,pc,leaf").input("test.rs").run();
+    // run("test");
 }
diff --git a/tests/ui/invalid-compile-flags/branch-protection-missing-pac-ret.BADFLAGS.stderr b/tests/ui/invalid-compile-flags/branch-protection-missing-pac-ret.BADFLAGS.stderr
index d0e8d4719d3..dae08119dbc 100644
--- a/tests/ui/invalid-compile-flags/branch-protection-missing-pac-ret.BADFLAGS.stderr
+++ b/tests/ui/invalid-compile-flags/branch-protection-missing-pac-ret.BADFLAGS.stderr
@@ -1,2 +1,2 @@
-error: incorrect value `leaf` for unstable option `branch-protection` - a `,` separated combination of `bti`, `b-key`, `pac-ret`, or `leaf` was expected
+error: incorrect value `leaf` for unstable option `branch-protection` - a `,` separated combination of `bti`, `pac-ret`, followed by a combination of `pc`, `b-key`, or `leaf` was expected
 
diff --git a/tests/ui/invalid-compile-flags/branch-protection-missing-pac-ret.BADFLAGSPC.stderr b/tests/ui/invalid-compile-flags/branch-protection-missing-pac-ret.BADFLAGSPC.stderr
new file mode 100644
index 00000000000..13f79e94674
--- /dev/null
+++ b/tests/ui/invalid-compile-flags/branch-protection-missing-pac-ret.BADFLAGSPC.stderr
@@ -0,0 +1,2 @@
+error: incorrect value `pc` for unstable option `branch-protection` - a `,` separated combination of `bti`, `pac-ret`, followed by a combination of `pc`, `b-key`, or `leaf` was expected
+
diff --git a/tests/ui/invalid-compile-flags/branch-protection-missing-pac-ret.rs b/tests/ui/invalid-compile-flags/branch-protection-missing-pac-ret.rs
index c0a4bcac11b..b4025080034 100644
--- a/tests/ui/invalid-compile-flags/branch-protection-missing-pac-ret.rs
+++ b/tests/ui/invalid-compile-flags/branch-protection-missing-pac-ret.rs
@@ -1,7 +1,10 @@
-//@ revisions: BADFLAGS BADTARGET
+//@ revisions: BADFLAGS BADFLAGSPC BADTARGET
 //@ [BADFLAGS] compile-flags: --target=aarch64-unknown-linux-gnu -Zbranch-protection=leaf
 //@ [BADFLAGS] check-fail
 //@ [BADFLAGS] needs-llvm-components: aarch64
+//@ [BADFLAGSPC] compile-flags: --target=aarch64-unknown-linux-gnu -Zbranch-protection=pc
+//@ [BADFLAGSPC] check-fail
+//@ [BADFLAGSPC] needs-llvm-components: aarch64
 //@ [BADTARGET] compile-flags: --target=x86_64-unknown-linux-gnu -Zbranch-protection=bti
 //@ [BADTARGET] check-fail
 //@ [BADTARGET] needs-llvm-components: x86
@@ -10,5 +13,5 @@
 #![feature(no_core, lang_items)]
 #![no_core]
 
-#[lang="sized"]
-trait Sized { }
+#[lang = "sized"]
+trait Sized {}