Allow section override when using patchable-function-entries

compiler/rustc_attr_parsing/src/attributes/codegen_attrs.rs

@@ -8,9 +8,9 @@ use rustc_span::edition::Edition::Edition2024;
 use super::prelude::*;
 use crate::attributes::AttributeSafety;
 use crate::session_diagnostics::{
-    EmptyExportName, NakedFunctionIncompatibleAttribute, NullOnExport, NullOnObjcClass,
-    NullOnObjcSelector, ObjcClassExpectedStringLiteral, ObjcSelectorExpectedStringLiteral,
-    SanitizeInvalidStatic,
+    EmptyExportName, EmptySection, NakedFunctionIncompatibleAttribute, NullOnExport,
+    NullOnObjcClass, NullOnObjcSelector, NullOnSection, ObjcClassExpectedStringLiteral,
+    ObjcSelectorExpectedStringLiteral, SanitizeInvalidStatic,
 };
 use crate::target_checking::Policy::AllowSilent;
 
@@ -781,82 +781,94 @@ pub(crate) struct PatchableFunctionEntryParser;
 impl SingleAttributeParser for PatchableFunctionEntryParser {
     const PATH: &[Symbol] = &[sym::patchable_function_entry];
     const ALLOWED_TARGETS: AllowedTargets<'_> = AllowedTargets::AllowList(&[Allow(Target::Fn)]);
-    const TEMPLATE: AttributeTemplate = template!(List: &["prefix_nops = m, entry_nops = n"]);
+    const TEMPLATE: AttributeTemplate =
+        template!(List: &["prefix_nops = m, entry_nops = n, section = \"section\""]);
     const STABILITY: AttributeStability = unstable!(patchable_function_entry);
 
     fn convert(cx: &mut AcceptContext<'_, '_>, args: &ArgParser) -> Option<AttributeKind> {
         let meta_item_list = cx.expect_list(args, cx.attr_span)?;
 
         let mut prefix = None;
         let mut entry = None;
+        let mut section = None;
 
         if meta_item_list.len() == 0 {
             cx.adcx().expected_at_least_one_argument(meta_item_list.span);
             return None;
         }
 
-        let mut errored = false;
-
         for item in meta_item_list.mixed() {
             let Some((ident, value)) = cx.expect_name_value(item, item.span(), None) else {
-                continue;
+                return None;
             };
 
             let attrib_to_write = match ident.name {
                 sym::prefix_nops => {
                     // Duplicate prefixes are not allowed
                     if prefix.is_some() {
-                        errored = true;
                         cx.adcx().duplicate_key(ident.span, sym::prefix_nops);
-                        continue;
+                        return None;
                     }
                     &mut prefix
                 }
                 sym::entry_nops => {
                     // Duplicate entries are not allowed
                     if entry.is_some() {
-                        errored = true;
                         cx.adcx().duplicate_key(ident.span, sym::entry_nops);
-                        continue;
+                        return None;
                     }
                     &mut entry
                 }
+                sym::section => {
+                    // Duplicate entries are not allowed
+                    if section.is_some() {
+                        cx.adcx().duplicate_key(ident.span, sym::section);
+                        return None;
+                    }
+                    // Only a string type value is allowed.
+                    let Some(value_str) = value.value_as_str() else {
+                        cx.adcx().expect_string_literal(value);
+                        return None;
+                    };
+                    // The section name does not allow null characters.
+                    if value_str.as_str().contains('\0') {
+                        cx.emit_err(NullOnSection { span: value.value_span });
+                    }
+                    // The section name is not allowed to be empty, LLVM does
+                    // not allow them.
+                    if value_str.is_empty() {
+                        cx.emit_err(EmptySection { span: value.value_span });
+                    }
+                    section = Some(value_str);
+                    // Integer parsing is not needed, process next item.
+                    continue;
+                }
                 _ => {
-                    errored = true;
                     cx.adcx().expected_specific_argument(
                         ident.span,
                         &[sym::prefix_nops, sym::entry_nops],
                     );
-                    continue;
+                    return None;
                 }
             };
 
             let rustc_ast::LitKind::Int(val, _) = value.value_as_lit().kind else {
-                errored = true;
                 cx.adcx().expected_integer_literal(value.value_span);
-                continue;
+                return None;
             };
 
             let Ok(val) = val.get().try_into() else {
-                errored = true;
                 cx.adcx().expected_integer_literal_in_range(
                     value.value_span,
                     u8::MIN as isize,
                     u8::MAX as isize,
                 );
-                continue;
+                return None;
             };
 
             *attrib_to_write = Some(val);
         }
 
-        if errored {
-            None
-        } else {
-            Some(AttributeKind::PatchableFunctionEntry {
-                prefix: prefix.unwrap_or(0),
-                entry: entry.unwrap_or(0),
-            })
-        }
+        Some(AttributeKind::PatchableFunctionEntry { prefix, entry, section })
     }
 }

compiler/rustc_attr_parsing/src/session_diagnostics.rs

@@ -265,6 +265,13 @@ pub(crate) struct EmptyExportName {
     pub span: Span,
 }
 
+#[derive(Diagnostic)]
+#[diag("`section` may not be empty")]
+pub(crate) struct EmptySection {
+    #[primary_span]
+    pub span: Span,
+}
+
 #[derive(Diagnostic)]
 #[diag("`export_name` may not contain null characters", code = E0648)]
 pub(crate) struct NullOnExport {
@@ -300,6 +307,13 @@ pub(crate) struct NullOnObjcSelector {
     pub span: Span,
 }
 
+#[derive(Diagnostic)]
+#[diag("`section` may not contain null characters", code = E0648)]
+pub(crate) struct NullOnSection {
+    #[primary_span]
+    pub span: Span,
+}
+
 #[derive(Diagnostic)]
 #[diag("`objc::class!` expected a string literal")]
 pub(crate) struct ObjcClassExpectedStringLiteral {

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -86,11 +86,26 @@ fn patchable_function_entry_attrs<'ll>(
     attr: Option<PatchableFunctionEntry>,
 ) -> SmallVec<[&'ll Attribute; 2]> {
     let mut attrs = SmallVec::new();
-    let patchable_spec = attr.unwrap_or_else(|| {
-        PatchableFunctionEntry::from_config(sess.opts.unstable_opts.patchable_function_entry)
-    });
-    let entry = patchable_spec.entry();
-    let prefix = patchable_spec.prefix();
+
+    let mut entry = sess.opts.unstable_opts.patchable_function_entry.entry();
+    let mut prefix = sess.opts.unstable_opts.patchable_function_entry.prefix();
+    let mut section = sess.opts.unstable_opts.patchable_function_entry.section();
+    let section_sym;
+
+    // Apply attribute specified overrides, if any.
+    if let Some(patchable_spec) = attr {
+        if let Some(sym) = patchable_spec.section() {
+            section_sym = sym;
+            section = Some(section_sym.as_str());
+        }
+        // Override the nop counts if either is present. If only one is present, the
+        // other count is implied to be 0.
+        if patchable_spec.entry().is_some() || patchable_spec.prefix().is_some() {
+            entry = patchable_spec.entry().unwrap_or(0);
+            prefix = patchable_spec.prefix().unwrap_or(0);
+        }
+    }
+
     if entry > 0 {
         attrs.push(llvm::CreateAttrStringValue(
             cx.llcx,
@@ -105,6 +120,13 @@ fn patchable_function_entry_attrs<'ll>(
             &format!("{}", prefix),
         ));
     }
+    if let Some(section) = section {
+        attrs.push(llvm::CreateAttrStringValue(
+            cx.llcx,
+            "patchable-function-entry-section",
+            section,
+        ));
+    }
     attrs
 }
 

compiler/rustc_codegen_llvm/src/context.rs

@@ -14,7 +14,6 @@ use rustc_data_structures::base_n::{ALPHANUMERIC_ONLY, ToBaseN};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::small_c_str::SmallCStr;
 use rustc_hir::def_id::DefId;
-use rustc_middle::middle::codegen_fn_attrs::PatchableFunctionEntry;
 use rustc_middle::mono::CodegenUnit;
 use rustc_middle::ty::layout::{
     FnAbiError, FnAbiOfHelpers, FnAbiRequest, HasTypingEnv, LayoutError, LayoutOfHelpers,
@@ -343,14 +342,13 @@ pub(crate) unsafe fn create_module<'ll>(
 
         // Add "kcfi-offset" module flag with -Z patchable-function-entry (See
         // https://reviews.llvm.org/D141172).
-        let pfe =
-            PatchableFunctionEntry::from_config(sess.opts.unstable_opts.patchable_function_entry);
-        if pfe.prefix() > 0 {
+        let patchable_prefix_nops = sess.opts.unstable_opts.patchable_function_entry.prefix();
+        if patchable_prefix_nops > 0 {
             llvm::add_module_flag_u32(
                 llmod,
                 llvm::ModuleFlagMergeBehavior::Override,
                 "kcfi-offset",
-                pfe.prefix().into(),
+                patchable_prefix_nops.into(),
             );
         }
 

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -290,9 +290,11 @@ fn process_builtin_attrs(
             AttributeKind::RustcOffloadKernel => {
                 codegen_fn_attrs.flags |= CodegenFnAttrFlags::OFFLOAD_KERNEL
             }
-            AttributeKind::PatchableFunctionEntry { prefix, entry } => {
+            AttributeKind::PatchableFunctionEntry { prefix, entry, section } => {
                 codegen_fn_attrs.patchable_function_entry =
-                    Some(PatchableFunctionEntry::from_prefix_and_entry(*prefix, *entry));
+                    Some(PatchableFunctionEntry::from_prefix_entry_and_section(
+                        *prefix, *entry, *section,
+                    ));
             }
             AttributeKind::InstrumentFn(instrument_fn) => {
                 codegen_fn_attrs.instrument_fn = match instrument_fn {

compiler/rustc_hir/src/attrs/data_structures.rs

@@ -1224,8 +1224,9 @@ pub enum AttributeKind {
 
     /// Represents `#[patchable_function_entry]`
     PatchableFunctionEntry {
-        prefix: u8,
-        entry: u8,
+        prefix: Option<u8>,
+        entry: Option<u8>,
+        section: Option<Symbol>,
     },
 
     /// Represents `#[path]`

compiler/rustc_interface/src/tests.rs

@@ -866,7 +866,7 @@ fn test_unstable_options_tracking_hash() {
     tracked!(panic_in_drop, PanicStrategy::Abort);
     tracked!(
         patchable_function_entry,
-        PatchableFunctionEntry::from_total_and_prefix_nops(10, 5)
+        PatchableFunctionEntry::from_parts(10, 5, None)
             .expect("total must be greater than or equal to prefix")
     );
     tracked!(plt, Some(true));

compiler/rustc_middle/src/middle/codegen_fn_attrs.rs

@@ -114,7 +114,7 @@ pub struct CodegenFnAttrs {
     // FIXME(#82232, #143834): temporarily renamed to mitigate `#[align]` nameres ambiguity
     pub alignment: Option<Align>,
     /// The `#[patchable_function_entry(...)]` attribute. Indicates how many nops should be around
-    /// the function entry.
+    /// the function entry, or override default section to record entry location.
     pub patchable_function_entry: Option<PatchableFunctionEntry>,
     /// The `#[rustc_objc_class = "..."]` attribute.
     pub objc_class: Option<Symbol>,
@@ -162,24 +162,33 @@ pub struct TargetFeature {
 #[derive(Copy, Clone, Debug, TyEncodable, TyDecodable, StableHash)]
 pub struct PatchableFunctionEntry {
     /// Nops to prepend to the function
-    prefix: u8,
+    prefix: Option<u8>,
     /// Nops after entry, but before body
-    entry: u8,
+    entry: Option<u8>,
+    /// Optional, specific section to record entry location in
+    section: Option<Symbol>,
 }
 
 impl PatchableFunctionEntry {
-    pub fn from_config(config: rustc_session::config::PatchableFunctionEntry) -> Self {
-        Self { prefix: config.prefix(), entry: config.entry() }
+    pub fn from_prefix_entry_and_section(
+        prefix: Option<u8>,
+        entry: Option<u8>,
+        section: Option<Symbol>,
+    ) -> Self {
+        Self { prefix, entry, section }
     }
     pub fn from_prefix_and_entry(prefix: u8, entry: u8) -> Self {
-        Self { prefix, entry }
+        Self { prefix: Some(prefix), entry: Some(entry), section: None }
     }
-    pub fn prefix(&self) -> u8 {
+    pub fn prefix(&self) -> Option<u8> {
         self.prefix
     }
-    pub fn entry(&self) -> u8 {
+    pub fn entry(&self) -> Option<u8> {
         self.entry
     }
+    pub fn section(&self) -> Option<Symbol> {
+        self.section
+    }
 }
 
 #[derive(Clone, Copy, PartialEq, Eq, TyEncodable, TyDecodable, StableHash)]

compiler/rustc_session/src/config.rs

@@ -3342,23 +3342,29 @@ impl DumpMonoStatsFormat {
 
 /// `-Z patchable-function-entry` representation - how many nops to put before and after function
 /// entry.
-#[derive(Clone, Copy, PartialEq, Hash, Debug, Default)]
+#[derive(Clone, PartialEq, Hash, Debug, Default)]
 pub struct PatchableFunctionEntry {
     /// Nops before the entry
     prefix: u8,
     /// Nops after the entry
     entry: u8,
+    /// An optional section name to record the entry location
+    section: Option<String>,
 }
 
 impl PatchableFunctionEntry {
-    pub fn from_total_and_prefix_nops(
+    pub fn from_parts(
         total_nops: u8,
         prefix_nops: u8,
+        section: Option<String>,
     ) -> Option<PatchableFunctionEntry> {
         if total_nops < prefix_nops {
             None
+        // Section name cannot contain null characters.
+        } else if section.as_ref().map(|x| x.contains('\0') || x.is_empty()).unwrap_or(false) {
+            None
         } else {
-            Some(Self { prefix: prefix_nops, entry: total_nops - prefix_nops })
+            Some(Self { prefix: prefix_nops, entry: total_nops - prefix_nops, section })
         }
     }
     pub fn prefix(&self) -> u8 {
@@ -3367,6 +3373,9 @@ impl PatchableFunctionEntry {
     pub fn entry(&self) -> u8 {
         self.entry
     }
+    pub fn section(&self) -> Option<&str> {
+        self.section.as_ref().map(|x| x.as_str())
+    }
 }
 
 /// `-Zpolonius` values, enabling the borrow checker polonius analysis, and which version: legacy,

compiler/rustc_session/src/options.rs

@@ -784,7 +784,7 @@ mod desc {
     pub(crate) const parse_passes: &str = "a space-separated list of passes, or `all`";
     pub(crate) const parse_panic_strategy: &str = "either `unwind`, `abort`, or `immediate-abort`";
     pub(crate) const parse_on_broken_pipe: &str = "either `kill`, `error`, or `inherit`";
-    pub(crate) const parse_patchable_function_entry: &str = "either two comma separated integers (total_nops,prefix_nops), with prefix_nops <= total_nops, or one integer (total_nops)";
+    pub(crate) const parse_patchable_function_entry: &str = "a comma separated list of (prefix_nops,total_nops,section_name), (prefix_nops,total_nops), or (total_nops). Where prefix_nops <= total_nops where 0 < total_nops <= 255 and prefix_nops <= total_nops";
     pub(crate) const parse_opt_panic_strategy: &str = parse_panic_strategy;
     pub(crate) const parse_relro_level: &str = "one of: `full`, `partial`, or `off`";
     pub(crate) const parse_sanitizers: &str = "comma separated list of sanitizers: `address`, `cfi`, `dataflow`, `hwaddress`, `kcfi`, `kernel-address`, `kernel-hwaddress`, `leak`, `memory`, `memtag`, `safestack`, `shadow-call-stack`, `thread`, or 'realtime'";
@@ -1206,20 +1206,24 @@ pub mod parse {
     ) -> bool {
         let mut total_nops = 0;
         let mut prefix_nops = 0;
+        let mut section = None;
 
         if !parse_number(&mut total_nops, v) {
-            let parts = v.and_then(|v| v.split_once(',')).unzip();
-            if !parse_number(&mut total_nops, parts.0) {
+            let parts: Vec<_> = v.unwrap_or("").split(',').collect();
+            if parts.len() < 2 || parts.len() > 3 {
                 return false;
             }
-            if !parse_number(&mut prefix_nops, parts.1) {
+
+            if !parse_number(&mut total_nops, Some(parts[0])) {
+                return false;
+            }
+            if !parse_number(&mut prefix_nops, Some(parts[1])) {
                 return false;
             }
+            section = parts.get(2).map(|x| x.to_string());
         }
 
-        if let Some(pfe) =
-            PatchableFunctionEntry::from_total_and_prefix_nops(total_nops, prefix_nops)
-        {
+        if let Some(pfe) = PatchableFunctionEntry::from_parts(total_nops, prefix_nops, section) {
             *slot = pfe;
             return true;
         }