Auto merge of #102963 - ilammy:xray-basic, r=estebank

Add `-Z instrument-xray` flag

Implement MCP https://github.com/rust-lang/compiler-team/issues/561, adding `-Z instrument-xray` flag which enables XRay instrumentation in LLVM.

4 files changed, 106 insertions, 3 deletions

diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs
index 973d860118e..7d2fdf94baa 100644
--- a/compiler/rustc_session/src/config.rs
+++ b/compiler/rustc_session/src/config.rs
@@ -174,6 +174,25 @@ pub enum InstrumentCoverage {
     Off,
 }
 
+/// Settings for `-Z instrument-xray` flag.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
+pub struct InstrumentXRay {
+    /// `-Z instrument-xray=always`, force instrumentation
+    pub always: bool,
+    /// `-Z instrument-xray=never`, disable instrumentation
+    pub never: bool,
+    /// `-Z instrument-xray=ignore-loops`, ignore presence of loops,
+    /// instrument functions based only on instruction count
+    pub ignore_loops: bool,
+    /// `-Z instrument-xray=instruction-threshold=N`, explicitly set instruction threshold
+    /// for instrumentation, or `None` to use compiler's default
+    pub instruction_threshold: Option<usize>,
+    /// `-Z instrument-xray=skip-entry`, do not instrument function entry
+    pub skip_entry: bool,
+    /// `-Z instrument-xray=skip-exit`, do not instrument function exit
+    pub skip_exit: bool,
+}
+
 #[derive(Clone, PartialEq, Hash, Debug)]
 pub enum LinkerPluginLto {
     LinkerPlugin(PathBuf),
@@ -2805,9 +2824,9 @@ impl PpMode {
 pub(crate) mod dep_tracking {
     use super::{
         BranchProtection, CFGuard, CFProtection, CrateType, DebugInfo, ErrorOutputType,
-        InstrumentCoverage, LdImpl, LinkerPluginLto, LocationDetail, LtoCli, OomStrategy, OptLevel,
-        OutputType, OutputTypes, Passes, SourceFileHashAlgorithm, SplitDwarfKind,
-        SwitchWithOptPath, SymbolManglingVersion, TraitSolver, TrimmedDefPaths,
+        InstrumentCoverage, InstrumentXRay, LdImpl, LinkerPluginLto, LocationDetail, LtoCli,
+        OomStrategy, OptLevel, OutputType, OutputTypes, Passes, SourceFileHashAlgorithm,
+        SplitDwarfKind, SwitchWithOptPath, SymbolManglingVersion, TraitSolver, TrimmedDefPaths,
     };
     use crate::lint;
     use crate::options::WasiExecModel;
@@ -2876,6 +2895,7 @@ pub(crate) mod dep_tracking {
         CodeModel,
         TlsModel,
         InstrumentCoverage,
+        InstrumentXRay,
         CrateType,
         MergeFunctions,
         PanicStrategy,
diff --git a/compiler/rustc_session/src/errors.rs b/compiler/rustc_session/src/errors.rs
index 8e8fba5e236..c851145440b 100644
--- a/compiler/rustc_session/src/errors.rs
+++ b/compiler/rustc_session/src/errors.rs
@@ -72,6 +72,12 @@ pub struct ProfileSampleUseFileDoesNotExist<'a> {
 pub struct TargetRequiresUnwindTables;
 
 #[derive(Diagnostic)]
+#[diag(session_instrumentation_not_supported)]
+pub struct InstrumentationNotSupported {
+    pub us: String,
+}
+
+#[derive(Diagnostic)]
 #[diag(session_sanitizer_not_supported)]
 pub struct SanitizerNotSupported {
     pub us: String,
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
index 61cb81aec3d..b9723d35e58 100644
--- a/compiler/rustc_session/src/options.rs
+++ b/compiler/rustc_session/src/options.rs
@@ -380,6 +380,7 @@ mod desc {
     pub const parse_dump_mono_stats: &str = "`markdown` (default) or `json`";
     pub const parse_instrument_coverage: &str =
         "`all` (default), `except-unused-generics`, `except-unused-functions`, or `off`";
+    pub const parse_instrument_xray: &str = "either a boolean (`yes`, `no`, `on`, `off`, etc), or a comma separated list of settings: `always` or `never` (mutually exclusive), `ignore-loops`, `instruction-threshold=N`, `skip-entry`, `skip-exit`";
     pub const parse_unpretty: &str = "`string` or `string=string`";
     pub const parse_treat_err_as_bug: &str = "either no value or a number bigger than 0";
     pub const parse_trait_solver: &str =
@@ -869,6 +870,68 @@ mod parse {
         true
     }
 
+    pub(crate) fn parse_instrument_xray(
+        slot: &mut Option<InstrumentXRay>,
+        v: Option<&str>,
+    ) -> bool {
+        if v.is_some() {
+            let mut bool_arg = None;
+            if parse_opt_bool(&mut bool_arg, v) {
+                *slot = if bool_arg.unwrap() { Some(InstrumentXRay::default()) } else { None };
+                return true;
+            }
+        }
+
+        let mut options = slot.get_or_insert_default();
+        let mut seen_always = false;
+        let mut seen_never = false;
+        let mut seen_ignore_loops = false;
+        let mut seen_instruction_threshold = false;
+        let mut seen_skip_entry = false;
+        let mut seen_skip_exit = false;
+        for option in v.into_iter().map(|v| v.split(',')).flatten() {
+            match option {
+                "always" if !seen_always && !seen_never => {
+                    options.always = true;
+                    options.never = false;
+                    seen_always = true;
+                }
+                "never" if !seen_never && !seen_always => {
+                    options.never = true;
+                    options.always = false;
+                    seen_never = true;
+                }
+                "ignore-loops" if !seen_ignore_loops => {
+                    options.ignore_loops = true;
+                    seen_ignore_loops = true;
+                }
+                option
+                    if option.starts_with("instruction-threshold")
+                        && !seen_instruction_threshold =>
+                {
+                    let Some(("instruction-threshold", n)) = option.split_once('=') else {
+                        return false;
+                    };
+                    match n.parse() {
+                        Ok(n) => options.instruction_threshold = Some(n),
+                        Err(_) => return false,
+                    }
+                    seen_instruction_threshold = true;
+                }
+                "skip-entry" if !seen_skip_entry => {
+                    options.skip_entry = true;
+                    seen_skip_entry = true;
+                }
+                "skip-exit" if !seen_skip_exit => {
+                    options.skip_exit = true;
+                    seen_skip_exit = true;
+                }
+                _ => return false,
+            }
+        }
+        true
+    }
+
     pub(crate) fn parse_treat_err_as_bug(slot: &mut Option<NonZeroUsize>, v: Option<&str>) -> bool {
         match v {
             Some(s) => {
@@ -1397,6 +1460,16 @@ options! {
         `=off` (default)"),
     instrument_mcount: bool = (false, parse_bool, [TRACKED],
         "insert function instrument code for mcount-based tracing (default: no)"),
+    instrument_xray: Option<InstrumentXRay> = (None, parse_instrument_xray, [TRACKED],
+        "insert function instrument code for XRay-based tracing (default: no)
+         Optional extra settings:
+         `=always`
+         `=never`
+         `=ignore-loops`
+         `=instruction-threshold=N`
+         `=skip-entry`
+         `=skip-exit`
+         Multiple options can be combined with commas."),
     keep_hygiene_data: bool = (false, parse_bool, [UNTRACKED],
         "keep hygiene data after analysis (default: no)"),
     layout_seed: Option<u64> = (None, parse_opt_number, [TRACKED],
diff --git a/compiler/rustc_session/src/session.rs b/compiler/rustc_session/src/session.rs
index 8a0176f6391..fe6ac80fde6 100644
--- a/compiler/rustc_session/src/session.rs
+++ b/compiler/rustc_session/src/session.rs
@@ -1589,6 +1589,10 @@ fn validate_commandline_args_with_session_available(sess: &Session) {
     {
         sess.emit_err(errors::SplitDebugInfoUnstablePlatform { debuginfo: sess.split_debuginfo() });
     }
+
+    if sess.opts.unstable_opts.instrument_xray.is_some() && !sess.target.options.supports_xray {
+        sess.emit_err(errors::InstrumentationNotSupported { us: "XRay".to_string() });
+    }
 }
 
 /// Holds data on the current incremental compilation session, if there is one.