Implement RFC 3917 for instrumented functions

compiler/rustc_attr_parsing/src/attributes/codegen_attrs.rs

@@ -586,6 +586,48 @@ impl<S: Stage> CombineAttributeParser<S> for ForceTargetFeatureParser {
     }
 }
 
+pub(crate) struct InstrumentFnParser;
+
+impl<S: Stage> SingleAttributeParser<S> 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<S> = OnDuplicate::Error;
+
+    fn convert(cx: &mut AcceptContext<'_, '_, S>, 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(_) => {
+                let span = cx.attr_span;
+                cx.adcx().expected_single_argument(span);
+            }
+            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<S: Stage> SingleAttributeParser<S> for SanitizeParser {

compiler/rustc_attr_parsing/src/context.rs

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

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -6,7 +6,9 @@ use rustc_middle::middle::codegen_fn_attrs::{
     CodegenFnAttrFlags, CodegenFnAttrs, PatchableFunctionEntry, SanitizerFnAttrs, TargetFeature,
 };
 use rustc_middle::ty::{self, TyCtxt};
-use rustc_session::config::{BranchProtection, FunctionReturn, OptLevel, PAuthKey, PacRet};
+use rustc_session::config::{
+    BranchProtection, FunctionReturn, InstrumentFunction, OptLevel, PAuthKey, PacRet,
+};
 use rustc_span::sym;
 use rustc_symbol_mangling::mangle_internal_symbol;
 use rustc_target::spec::{Arch, FramePointer, SanitizerSet, StackProbeType, StackProtector};
@@ -175,7 +177,7 @@ pub(crate) fn frame_pointer_type_attr<'ll>(
     let opts = &sess.opts;
     // "mcount" function relies on stack pointer.
     // See <https://sourceware.org/binutils/docs/gprof/Implementation.html>.
-    if opts.unstable_opts.instrument_mcount {
+    if opts.unstable_opts.instrument_function == InstrumentFunction::Mcount {
         fp.ratchet(FramePointer::Always);
     }
     fp.ratchet(opts.cg.force_frame_pointers);
@@ -201,52 +203,104 @@ 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.
+    match sess.opts.unstable_opts.instrument_function {
+        InstrumentFunction::Fentry => {
+            // Similar to `clang -pg -mfentry` behavior.
 
-        // 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 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"));
-        }
-        if options.never {
-            attrs.push(llvm::CreateAttrStringValue(cx.llcx, "function-instrument", "xray-never"));
-        }
-        if options.ignore_loops {
-            attrs.push(llvm::CreateAttrString(cx.llcx, "xray-ignore-loops"));
+            if instrument_entry {
+                attrs.push(llvm::CreateAttrStringValue(cx.llcx, "fentry-call", "true"));
+
+                if sess.opts.unstable_opts.instrument_fentry_opts.no_call {
+                    attrs.push(llvm::CreateAttrString(cx.llcx, "mnop-mcount"));
+                }
+                if sess.opts.unstable_opts.instrument_fentry_opts.record {
+                    attrs.push(llvm::CreateAttrString(cx.llcx, "mrecord-mcount"));
+                }
+            }
         }
-        // LLVM will not choose the default for us, but rather requires specific
-        // threshold in absence of "xray-always". Use the same default as Clang.
-        let threshold = options.instruction_threshold.unwrap_or(200);
-        attrs.push(llvm::CreateAttrStringValue(
-            cx.llcx,
-            "xray-instruction-threshold",
-            &threshold.to_string(),
-        ));
-        if options.skip_entry {
-            attrs.push(llvm::CreateAttrString(cx.llcx, "xray-skip-entry"));
+        InstrumentFunction::Mcount => {
+            // Similar to `clang -pg` behavior. Handled by the
+            // `post-inline-ee-instrument` LLVM pass.
+
+            // #[instrument_fn], the default is on.
+            let instrument_entry = instrument_fn.unwrap_or_else(|| true);
+
+            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 options.skip_exit {
-            attrs.push(llvm::CreateAttrString(cx.llcx, "xray-skip-exit"));
+        InstrumentFunction::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.
+            let options = &sess.opts.unstable_opts.instrument_xray_opts;
+
+            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 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"));
+            }
+            // LLVM will not choose the default for us, but rather requires specific
+            // threshold in absence of "xray-always". Use the same default as Clang.
+            let threshold = options.instruction_threshold.unwrap_or(200);
+            attrs.push(llvm::CreateAttrStringValue(
+                cx.llcx,
+                "xray-instruction-threshold",
+                &threshold.to_string(),
+            ));
+
+            if options.skip_entry {
+                attrs.push(llvm::CreateAttrString(cx.llcx, "xray-skip-entry"));
+            }
+            if options.skip_exit {
+                attrs.push(llvm::CreateAttrString(cx.llcx, "xray-skip-exit"));
+            }
         }
+        InstrumentFunction::No => {}
     }
     attrs
 }
@@ -429,7 +483,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/back/link.rs

@@ -35,8 +35,8 @@ use rustc_middle::middle::debugger_visualizer::DebuggerVisualizerFile;
 use rustc_middle::middle::dependency_format::Linkage;
 use rustc_middle::middle::exported_symbols::SymbolExportKind;
 use rustc_session::config::{
-    self, CFGuard, CrateType, DebugInfo, LinkerFeaturesCli, OutFileName, OutputFilenames,
-    OutputType, PrintKind, SplitDwarfKind, Strip,
+    self, CFGuard, CrateType, DebugInfo, InstrumentFunction, LinkerFeaturesCli, OutFileName,
+    OutputFilenames, OutputType, PrintKind, SplitDwarfKind, Strip,
 };
 use rustc_session::lint::builtin::LINKER_MESSAGES;
 use rustc_session::output::{check_file_is_writeable, invalid_output_for_target, out_filename};
@@ -2767,8 +2767,11 @@ fn add_order_independent_options(
         cmd.pgo_gen();
     }
 
-    if sess.opts.unstable_opts.instrument_mcount {
-        cmd.enable_profiling();
+    match sess.opts.unstable_opts.instrument_function {
+        InstrumentFunction::Mcount | InstrumentFunction::Fentry => {
+            cmd.enable_profiling();
+        }
+        _ => {}
     }
 
     if sess.opts.cg.control_flow_guard != CFGuard::Disabled {

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

@@ -503,6 +503,13 @@ pub static BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         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

@@ -550,6 +550,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(99999)),
     /// 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

@@ -1209,6 +1209,9 @@ pub enum AttributeKind {
     /// Represents `#[instruction_set]`
     InstructionSet(InstructionSetAttr),
 
+    /// Represents `#[instrument_fn]`
+    InstrumentFn(Option<bool>),
+
     /// Represents `#[lang]`
     Lang(LangItem, Span),
 

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_interface/src/tests.rs

@@ -13,11 +13,12 @@ use rustc_session::config::{
     AnnotateMoves, AutoDiff, BranchProtection, CFGuard, Cfg, CoverageLevel, CoverageOptions,
     DebugInfo, DumpMonoStatsFormat, ErrorOutputType, ExternEntry, ExternLocation, Externs,
     FmtDebug, FunctionReturn, IncrementalStateAssertion, InliningThreshold, Input,
-    InstrumentCoverage, InstrumentXRay, LinkSelfContained, LinkerPluginLto, LocationDetail, LtoCli,
-    MirIncludeSpans, NextSolverConfig, Offload, Options, OutFileName, OutputType, OutputTypes,
-    PAuthKey, PacRet, Passes, PatchableFunctionEntry, Polonius, ProcMacroExecutionStrategy, Strip,
-    SwitchWithOptPath, SymbolManglingVersion, WasiExecModel, build_configuration,
-    build_session_options, rustc_optgroups,
+    InstrumentCoverage, InstrumentFunction, InstrumentMcountOpts, InstrumentXRayOpts,
+    LinkSelfContained, LinkerPluginLto, LocationDetail, LtoCli, MirIncludeSpans, NextSolverConfig,
+    Offload, Options, OutFileName, OutputType, OutputTypes, PAuthKey, PacRet, Passes,
+    PatchableFunctionEntry, Polonius, ProcMacroExecutionStrategy, Strip, SwitchWithOptPath,
+    SymbolManglingVersion, WasiExecModel, build_configuration, build_session_options,
+    rustc_optgroups,
 };
 use rustc_session::lint::Level;
 use rustc_session::search_paths::SearchPath;
@@ -810,8 +811,12 @@ fn test_unstable_options_tracking_hash() {
     tracked!(inline_mir, Some(true));
     tracked!(inline_mir_hint_threshold, Some(123));
     tracked!(inline_mir_threshold, Some(123));
-    tracked!(instrument_mcount, true);
-    tracked!(instrument_xray, Some(InstrumentXRay::default()));
+    tracked!(instrument_fentry_opts, InstrumentMcountOpts { no_call: true, record: true });
+    tracked!(instrument_function, InstrumentFunction::Fentry);
+    tracked!(
+        instrument_xray_opts,
+        InstrumentXRayOpts { always: true, ..InstrumentXRayOpts::default() }
+    );
     tracked!(link_directives, false);
     tracked!(link_only, true);
     tracked!(lint_llvm_ir, true);

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, HashStable, 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

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