Add instrument_fn attribute

This attribute enables or disables function instrumentation
when using `-Zinstrument-xray` or `-Zinstrument-mcount`.

It supports the following usage:

`#[instrument_fn = "on|off"]`

For XRay, "on" is equivalent to always instrument, and "off" is
equivalent to never instrumenting.

For mcount, "on" has no effect. "off" disables instrumentation
of the function.

Author: pmur

compiler/rustc_attr_parsing/src/attributes/codegen_attrs.rs

@@ -1,5 +1,7 @@
 use rustc_feature::AttributeStability;
-use rustc_hir::attrs::{CoverageAttrKind, OptimizeAttr, RtsanSetting, SanitizerSet, UsedBy};
+use rustc_hir::attrs::{
+    CoverageAttrKind, InstrumentFnAttr, OptimizeAttr, RtsanSetting, SanitizerSet, UsedBy,
+};
 use rustc_session::errors::feature_err;
 use rustc_span::edition::Edition::Edition2024;
 
@@ -563,6 +565,44 @@ impl CombineAttributeParser for ForceTargetFeatureParser {
     }
 }
 
+pub(crate) struct InstrumentFnParser;
+
+impl SingleAttributeParser for InstrumentFnParser {
+    const PATH: &[Symbol] = &[sym::instrument_fn];
+    const ALLOWED_TARGETS: AllowedTargets = AllowedTargets::AllowList(&[
+        Allow(Target::Fn),
+        Allow(Target::Method(MethodKind::Inherent)),
+        Allow(Target::Method(MethodKind::Trait { body: true })),
+        Allow(Target::Method(MethodKind::TraitImpl)),
+    ]);
+    const TEMPLATE: AttributeTemplate = template!(NameValueStr: "on|off");
+    const STABILITY: AttributeStability = unstable!(instrument_fn);
+
+    fn convert(cx: &mut AcceptContext<'_, '_>, args: &ArgParser) -> Option<AttributeKind> {
+        let instrument = match args {
+            ArgParser::NameValue(nv) => match nv.value_as_str() {
+                Some(sym::on) => Some(AttributeKind::InstrumentFn(InstrumentFnAttr::On)),
+                Some(sym::off) => Some(AttributeKind::InstrumentFn(InstrumentFnAttr::Off)),
+                _ => {
+                    cx.adcx()
+                        .expected_specific_argument_strings(nv.value_span, &[sym::on, sym::off]);
+                    None
+                }
+            },
+            ArgParser::List(l) => {
+                cx.adcx().expected_single_argument(l.span, l.len());
+                None
+            }
+            ArgParser::NoArgs => {
+                let span = cx.attr_span;
+                cx.adcx().expected_specific_argument_strings(span, &[sym::on, sym::off]);
+                None
+            }
+        };
+        instrument
+    }
+}
+
 pub(crate) struct SanitizeParser;
 
 impl SingleAttributeParser for SanitizeParser {

compiler/rustc_attr_parsing/src/context.rs

@@ -185,6 +185,7 @@ attribute_parsers!(
         Single<IgnoreParser>,
         Single<InlineParser>,
         Single<InstructionSetParser>,
+        Single<InstrumentFnParser>,
         Single<LangParser>,
         Single<LinkNameParser>,
         Single<LinkOrdinalParser>,

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -3,7 +3,8 @@ use rustc_hir::attrs::{InlineAttr, InstructionSetAttr, OptimizeAttr, RtsanSettin
 use rustc_hir::def_id::DefId;
 use rustc_hir::find_attr;
 use rustc_middle::middle::codegen_fn_attrs::{
-    CodegenFnAttrFlags, CodegenFnAttrs, PatchableFunctionEntry, SanitizerFnAttrs, TargetFeature,
+    CodegenFnAttrFlags, CodegenFnAttrs, InstrumentFnAttr, PatchableFunctionEntry, SanitizerFnAttrs,
+    TargetFeature,
 };
 use rustc_middle::ty::{self, Instance, TyCtxt};
 use rustc_session::config::{BranchProtection, FunctionReturn, OptLevel, PAuthKey, PacRet};
@@ -210,35 +211,59 @@ fn function_return_attr<'ll>(cx: &SimpleCx<'ll>, sess: &Session) -> Option<&'ll
 fn instrument_function_attr<'ll>(
     cx: &SimpleCx<'ll>,
     sess: &Session,
+    instrument_fn: InstrumentFnAttr,
 ) -> SmallVec<[&'ll Attribute; 4]> {
     let mut attrs = SmallVec::new();
     if sess.opts.unstable_opts.instrument_mcount {
         // Similar to `clang -pg` behavior. Handled by the
         // `post-inline-ee-instrument` LLVM pass.
 
-        // The function name varies on platforms.
-        // See test/CodeGen/mcount.c in clang.
-        let mcount_name = match &sess.target.llvm_mcount_intrinsic {
-            Some(llvm_mcount_intrinsic) => llvm_mcount_intrinsic.as_ref(),
-            None => sess.target.mcount.as_ref(),
+        let instrument_entry = match instrument_fn {
+            InstrumentFnAttr::Default | InstrumentFnAttr::On => true,
+            InstrumentFnAttr::Off => false,
         };
 
-        attrs.push(llvm::CreateAttrStringValue(
-            cx.llcx,
-            "instrument-function-entry-inlined",
-            mcount_name,
-        ));
+        if instrument_entry {
+            // The function name varies on platforms.
+            // See test/CodeGen/mcount.c in clang.
+            let mcount_name = match &sess.target.llvm_mcount_intrinsic {
+                Some(llvm_mcount_intrinsic) => llvm_mcount_intrinsic.as_ref(),
+                None => sess.target.mcount.as_ref(),
+            };
+
+            attrs.push(llvm::CreateAttrStringValue(
+                cx.llcx,
+                "instrument-function-entry-inlined",
+                mcount_name,
+            ));
+        }
     }
     if let Some(options) = &sess.opts.unstable_opts.instrument_xray {
         // XRay instrumentation is similar to __cyg_profile_func_{enter,exit}.
         // Function prologue and epilogue are instrumented with NOP sleds,
         // a runtime library later replaces them with detours into tracing code.
-        if options.always {
-            attrs.push(llvm::CreateAttrStringValue(cx.llcx, "function-instrument", "xray-always"));
+
+        let mut never = options.never;
+        let mut always = options.always;
+
+        // always and never may be overridden by the #[instrument_fn = ...] attribute.
+        match instrument_fn {
+            InstrumentFnAttr::Default => {}
+            InstrumentFnAttr::On => {
+                always = true;
+            }
+            InstrumentFnAttr::Off => {
+                never = true;
+            }
         }
-        if options.never {
+
+        if never {
             attrs.push(llvm::CreateAttrStringValue(cx.llcx, "function-instrument", "xray-never"));
         }
+        if always {
+            attrs.push(llvm::CreateAttrStringValue(cx.llcx, "function-instrument", "xray-always"));
+        }
+
         if options.ignore_loops {
             attrs.push(llvm::CreateAttrString(cx.llcx, "xray-ignore-loops"));
         }
@@ -442,7 +467,7 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
     // FIXME: none of these functions interact with source level attributes.
     to_add.extend(frame_pointer_type_attr(cx, sess));
     to_add.extend(function_return_attr(cx, sess));
-    to_add.extend(instrument_function_attr(cx, sess));
+    to_add.extend(instrument_function_attr(cx, sess, codegen_fn_attrs.instrument_fn));
     to_add.extend(nojumptables_attr(cx, sess));
     to_add.extend(probestack_attr(cx, tcx));
     to_add.extend(stackprotector_attr(cx, sess));

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -1,13 +1,14 @@
 use rustc_abi::{Align, ExternAbi};
 use rustc_hir::attrs::{
-    AttributeKind, EiiImplResolution, InlineAttr, Linkage, RtsanSetting, UsedBy,
+    AttributeKind, EiiImplResolution, InlineAttr, InstrumentFnAttr as HirInstrumentFnAttr, Linkage,
+    RtsanSetting, UsedBy,
 };
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LOCAL_CRATE, LocalDefId};
 use rustc_hir::{self as hir, Attribute, find_attr};
 use rustc_macros::Diagnostic;
 use rustc_middle::middle::codegen_fn_attrs::{
-    CodegenFnAttrFlags, CodegenFnAttrs, PatchableFunctionEntry, SanitizerFnAttrs,
+    CodegenFnAttrFlags, CodegenFnAttrs, InstrumentFnAttr, PatchableFunctionEntry, SanitizerFnAttrs,
 };
 use rustc_middle::mono::Visibility;
 use rustc_middle::query::Providers;
@@ -293,6 +294,12 @@ fn process_builtin_attrs(
                 codegen_fn_attrs.patchable_function_entry =
                     Some(PatchableFunctionEntry::from_prefix_and_entry(*prefix, *entry));
             }
+            AttributeKind::InstrumentFn(instrument_fn) => {
+                codegen_fn_attrs.instrument_fn = match instrument_fn {
+                    HirInstrumentFnAttr::On => InstrumentFnAttr::On,
+                    HirInstrumentFnAttr::Off => InstrumentFnAttr::Off,
+                };
+            }
             _ => {}
         }
     }

compiler/rustc_feature/src/builtin_attrs.rs

@@ -208,6 +208,11 @@ pub static BUILTIN_ATTRIBUTES: &[Symbol] = &[
     // - https://github.com/rust-lang/rust/issues/130494
     sym::pin_v2,
 
+    // RFC 3917
+    // `#[instrument_fn = "on|off"]` to insert or inhibit instrumentation function
+    // calls inside a function, usually around the prologue.
+    sym::instrument_fn,
+
     // ==========================================================================
     // Internal attributes: Stability, deprecation, and unsafe:
     // ==========================================================================

compiler/rustc_feature/src/unstable.rs

@@ -593,6 +593,8 @@ declare_features! (
     (unstable, import_trait_associated_functions, "1.86.0", Some(134691)),
     /// Allows associated types in inherent impls.
     (incomplete, inherent_associated_types, "1.52.0", Some(8995)),
+    /// Enable #[instrument_fn] on function.
+    (unstable, instrument_fn, "CURRENT_RUSTC_VERSION", Some(157081)),
     /// Allows using `pointer` and `reference` in intra-doc links
     (unstable, intra_doc_pointers, "1.51.0", Some(80896)),
     /// lahfsahf target feature on x86.

compiler/rustc_hir/src/attrs/data_structures.rs

@@ -135,6 +135,14 @@ pub enum InstructionSetAttr {
     ArmT32,
 }
 
+#[derive(Copy, Clone, PartialEq, Encodable, Decodable, Debug, Eq, StableHash, PrintAttribute)]
+pub enum InstrumentFnAttr {
+    /// `#[instrument_fn = "on"]`
+    On,
+    /// `#[instrument_fn = "off"]`
+    Off,
+}
+
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Default, PrintAttribute)]
 #[derive(Encodable, Decodable, StableHash)]
 pub enum OptimizeAttr {
@@ -1060,6 +1068,9 @@ pub enum AttributeKind {
     /// Represents `#[instruction_set]`
     InstructionSet(InstructionSetAttr),
 
+    /// Represents `#[instrument_fn]`
+    InstrumentFn(InstrumentFnAttr),
+
     /// Represents `#[lang]`
     Lang(LangItem),
 

compiler/rustc_hir/src/attrs/encode_cross_crate.rs

@@ -50,6 +50,7 @@ impl AttributeKind {
             Ignore { .. } => No,
             Inline(..) => No,
             InstructionSet(..) => No,
+            InstrumentFn(..) => No,
             Lang(..) => Yes,
             Link(..) => No,
             LinkName { .. } => Yes, // Needed for rustdoc

compiler/rustc_middle/src/middle/codegen_fn_attrs.rs

@@ -109,6 +109,24 @@ pub struct CodegenFnAttrs {
     pub objc_class: Option<Symbol>,
     /// The `#[rustc_objc_selector = "..."]` attribute.
     pub objc_selector: Option<Symbol>,
+    /// The `#[instrument_fn]` attribute.
+    pub instrument_fn: InstrumentFnAttr,
+}
+
+#[derive(Copy, Clone, TyEncodable, TyDecodable, StableHash, Debug)]
+pub enum InstrumentFnAttr {
+    /// Always instrument function
+    On,
+    /// Never instrument function
+    Off,
+    /// Instrument based on command line options, if any.
+    Default,
+}
+
+const impl Default for InstrumentFnAttr {
+    fn default() -> Self {
+        InstrumentFnAttr::Default
+    }
 }
 
 #[derive(Copy, Clone, Debug, TyEncodable, TyDecodable, StableHash, PartialEq, Eq)]
@@ -236,6 +254,7 @@ impl CodegenFnAttrs {
             patchable_function_entry: None,
             objc_class: None,
             objc_selector: None,
+            instrument_fn: InstrumentFnAttr::default(),
         }
     }
 

compiler/rustc_passes/src/check_attr.rs

@@ -277,6 +277,7 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
             AttributeKind::Fundamental => (),
             AttributeKind::Ignore { .. } => (),
             AttributeKind::InstructionSet(..) => (),
+            AttributeKind::InstrumentFn(..) => (),
             AttributeKind::Lang(..) => (),
             AttributeKind::LinkName { .. } => (),
             AttributeKind::LinkOrdinal { .. } => (),