Parse "-Z instrument-xray" codegen option

Recognize all bells and whistles that LLVM's XRay pass is capable of.
The always/never settings are a bit dumb without attributes but they're
still there. The default instruction count is chosen by the compiler,
not LLVM pass. We'll do it later.

5 files changed, 146 insertions, 3 deletions

diff --git a/compiler/rustc_interface/src/tests.rs b/compiler/rustc_interface/src/tests.rs
index 52a4e0e7418..5daefadabba 100644
--- a/compiler/rustc_interface/src/tests.rs
+++ b/compiler/rustc_interface/src/tests.rs
@@ -5,6 +5,7 @@ use rustc_data_structures::fx::FxHashSet;
 use rustc_errors::{emitter::HumanReadableErrorType, registry, ColorConfig};
 use rustc_session::config::rustc_optgroups;
 use rustc_session::config::Input;
+use rustc_session::config::InstrumentXRay;
 use rustc_session::config::TraitSolver;
 use rustc_session::config::{build_configuration, build_session_options, to_crate_config};
 use rustc_session::config::{
@@ -755,6 +756,7 @@ fn test_unstable_options_tracking_hash() {
     tracked!(inline_mir_threshold, Some(123));
     tracked!(instrument_coverage, Some(InstrumentCoverage::All));
     tracked!(instrument_mcount, true);
+    tracked!(instrument_xray, Some(InstrumentXRay::default()));
     tracked!(link_only, true);
     tracked!(llvm_plugins, vec![String::from("plugin_name")]);
     tracked!(location_detail, LocationDetail { file: true, line: false, column: false });
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/options.rs b/compiler/rustc_session/src/options.rs
index 0db4d85ff4b..a0a8a42575e 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/src/doc/unstable-book/src/compiler-flags/instrument-xray.md b/src/doc/unstable-book/src/compiler-flags/instrument-xray.md
new file mode 100644
index 00000000000..7fb33cd68b4
--- /dev/null
+++ b/src/doc/unstable-book/src/compiler-flags/instrument-xray.md
@@ -0,0 +1,39 @@
+# `instrument-xray`
+
+The tracking issue for this feature is: [#102921](https://github.com/rust-lang/rust/issues/102921).
+
+------------------------
+
+Enable generation of NOP sleds for XRay function tracing instrumentation.
+For more information on XRay,
+read [LLVM documentation](https://llvm.org/docs/XRay.html),
+and/or the [XRay whitepaper](http://research.google.com/pubs/pub45287.html).
+
+Set the `-Z instrument-xray` compiler flag in order to enable XRay instrumentation.
+
+  - `-Z instrument-xray` – use the default settings
+  - `-Z instrument-xray=skip-exit` – configure a custom setting
+  - `-Z instrument-xray=ignore-loops,instruction-threshold=300` –
+    multiple settings separated by commas
+
+Supported options:
+
+  - `always` – force instrumentation of all functions
+  - `never` – do no instrument any functions
+  - `ignore-loops` – ignore presence of loops,
+    instrument functions based only on instruction count
+  - `instruction-threshold=10` – set a different instruction threshold for instrumentation
+  - `skip-entry` – do no instrument function entry
+  - `skip-exit` – do no instrument function exit
+
+The default settings are:
+
+  - instrument both entry & exit from functions
+  - instrument functions with at least 200 instructions,
+    or containing a non-trivial loop
+
+Note that `-Z instrument-xray` only enables generation of NOP sleds
+which on their own don't do anything useful.
+In order to actually trace the functions,
+you will need to link a separate runtime library of your choice,
+such as Clang's [XRay Runtime Library](https://www.llvm.org/docs/XRay.html#xray-runtime-library).
diff --git a/tests/rustdoc-ui/z-help.stdout b/tests/rustdoc-ui/z-help.stdout
index 4f07fca82d1..2169b89c929 100644
--- a/tests/rustdoc-ui/z-help.stdout
+++ b/tests/rustdoc-ui/z-help.stdout
@@ -70,6 +70,15 @@
         `=except-unused-functions`
         `=off` (default)
     -Z                       instrument-mcount=val -- insert function instrument code for mcount-based tracing (default: no)
+    -Z                         instrument-xray=val -- 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.
     -Z                       keep-hygiene-data=val -- keep hygiene data after analysis (default: no)
     -Z                             layout-seed=val -- seed layout randomization
     -Z                   link-native-libraries=val -- link native libraries in the linker invocation (default: yes)