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

@@ -550,6 +550,47 @@ 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 ON_DUPLICATE: OnDuplicate = OnDuplicate::Error;
+
+    fn convert(cx: &mut AcceptContext<'_, '_>, args: &ArgParser) -> Option<AttributeKind> {
+        let mut instrument = None;
+        match args {
+            ArgParser::NameValue(nv) => {
+                if let Some(option) = nv.value_as_str()
+                    && (option == sym::off || option == sym::on)
+                {
+                    instrument = Some(option == sym::on);
+                } else {
+                    cx.adcx()
+                        .expected_specific_argument_strings(nv.value_span, &[sym::on, sym::off]);
+                }
+            }
+            ArgParser::List(l) => {
+                cx.adcx().expected_single_argument(l.span, l.len());
+            }
+            ArgParser::NoArgs => {
+                let span = cx.attr_span;
+                cx.adcx().expected_specific_argument_strings(span, &[sym::on, sym::off]);
+            }
+        }
+        Some(AttributeKind::InstrumentFn(instrument))
+    }
+}
+
 pub(crate) struct SanitizeParser;
 
 impl SingleAttributeParser for SanitizeParser {

compiler/rustc_attr_parsing/src/context.rs

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

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -205,35 +205,57 @@ fn function_return_attr<'ll>(cx: &SimpleCx<'ll>, sess: &Session) -> Option<&'ll
 fn instrument_function_attr<'ll>(
     cx: &SimpleCx<'ll>,
     sess: &Session,
+    instrument_fn: &Option<bool>,
 ) -> 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(),
-        };
+        // #[instrument_fn], the default is on.
+        let instrument_entry = instrument_fn.unwrap_or_else(|| true);
 
-        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;
+
+        // Apply optional #[instrument_fn] override.
+        match instrument_fn {
+            Some(true) => {
+                always = true;
+            }
+            Some(false) => {
+                never = true;
+            }
+            None => {}
         }
-        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"));
         }
@@ -245,6 +267,7 @@ fn instrument_function_attr<'ll>(
             "xray-instruction-threshold",
             &threshold.to_string(),
         ));
+
         if options.skip_entry {
             attrs.push(llvm::CreateAttrString(cx.llcx, "xray-skip-entry"));
         }
@@ -437,7 +460,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

@@ -293,6 +293,9 @@ 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 = *instrument_fn;
+            }
             _ => {}
         }
     }

compiler/rustc_feature/src/builtin_attrs.rs

@@ -401,6 +401,13 @@ pub static BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
     // - https://github.com/rust-lang/rust/issues/130494
     gated!(pin_v2, pin_ergonomics, experimental!(pin_v2)),
 
+    // RFC 3917
+    // `#[instrument_fn = "on|off"]` to insert or inhibit instrumentation function
+    // calls inside a function, usually around the prologue.
+    gated!(
+        instrument_fn, instrument_fn, experimental!(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 (todo: tracking issue)
+    (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

@@ -1060,6 +1060,9 @@ pub enum AttributeKind {
     /// Represents `#[instruction_set]`
     InstructionSet(InstructionSetAttr),
 
+    /// Represents `#[instrument_fn]`
+    InstrumentFn(Option<bool>),
+
     /// 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,8 @@ 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: Option<bool>,
 }
 
 #[derive(Copy, Clone, Debug, TyEncodable, TyDecodable, StableHash, PartialEq, Eq)]
@@ -236,6 +238,7 @@ impl CodegenFnAttrs {
             patchable_function_entry: None,
             objc_class: None,
             objc_selector: None,
+            instrument_fn: None,
         }
     }
 

compiler/rustc_passes/src/check_attr.rs

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