Rollup merge of #157681 - pmur:murp/add-instrument_fn-attr, r=oli-obk,mejrs

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.

This is tracked by #157081, and related to the pending rfc rust-lang/rfcs#3917, and earlier as part mcp rust-lang/compiler-team#561

Author: jhpratt

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;
 
@@ -575,6 +577,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

@@ -189,6 +189,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

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

compiler/rustc_feature/src/unstable.rs

@@ -595,6 +595,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

@@ -143,6 +143,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 {
@@ -1068,6 +1076,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

@@ -117,6 +117,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)]
@@ -244,6 +262,7 @@ impl CodegenFnAttrs {
             patchable_function_entry: None,
             objc_class: None,
             objc_selector: None,
+            instrument_fn: InstrumentFnAttr::default(),
         }
     }
 

compiler/rustc_passes/src/check_attr.rs

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