Add `-Zfunction-return={keep,thunk-extern}` option

This is intended to be used for Linux kernel RETHUNK builds.

With this commit (optionally backported to Rust 1.73.0), plus a
patched Linux kernel to pass the flag, I get a RETHUNK build with
Rust enabled that is `objtool`-warning-free and is able to boot in
QEMU and load a sample Rust kernel module.

Signed-off-by: Miguel Ojeda <ojeda@kernel.org>

4 files changed, 58 insertions, 4 deletions

diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs
index d6e7afb7d09..77874bd9c07 100644
--- a/compiler/rustc_session/src/config.rs
+++ b/compiler/rustc_session/src/config.rs
@@ -3164,9 +3164,9 @@ impl PpMode {
 pub(crate) mod dep_tracking {
     use super::{
         BranchProtection, CFGuard, CFProtection, CrateType, DebugInfo, DebugInfoCompression,
-        ErrorOutputType, InliningThreshold, InstrumentCoverage, InstrumentXRay, LinkerPluginLto,
-        LocationDetail, LtoCli, OomStrategy, OptLevel, OutFileName, OutputType, OutputTypes,
-        Polonius, RemapPathScopeComponents, ResolveDocLinks, SourceFileHashAlgorithm,
+        ErrorOutputType, FunctionReturn, InliningThreshold, InstrumentCoverage, InstrumentXRay,
+        LinkerPluginLto, LocationDetail, LtoCli, OomStrategy, OptLevel, OutFileName, OutputType,
+        OutputTypes, Polonius, RemapPathScopeComponents, ResolveDocLinks, SourceFileHashAlgorithm,
         SplitDwarfKind, SwitchWithOptPath, SymbolManglingVersion, TraitSolver, TrimmedDefPaths,
     };
     use crate::lint;
@@ -3273,6 +3273,7 @@ pub(crate) mod dep_tracking {
         TraitSolver,
         Polonius,
         InliningThreshold,
+        FunctionReturn,
     );
 
     impl<T1, T2> DepTrackingHash for (T1, T2)
@@ -3451,3 +3452,14 @@ impl Default for InliningThreshold {
         Self::Sometimes(100)
     }
 }
+
+/// The different settings that the `-Zfunction-return` flag can have.
+#[derive(Clone, Copy, PartialEq, Hash, Debug, Default)]
+pub enum FunctionReturn {
+    /// Keep the function return unmodified.
+    #[default]
+    Keep,
+
+    /// Replace returns with jumps to thunk, without emitting the thunk.
+    ThunkExtern,
+}
diff --git a/compiler/rustc_session/src/errors.rs b/compiler/rustc_session/src/errors.rs
index 31094e0d266..70ee46ea902 100644
--- a/compiler/rustc_session/src/errors.rs
+++ b/compiler/rustc_session/src/errors.rs
@@ -436,3 +436,11 @@ pub struct IncompatibleLinkerFlavor {
     pub flavor: &'static str,
     pub compatible_list: String,
 }
+
+#[derive(Diagnostic)]
+#[diag(session_function_return_requires_x86_or_x86_64)]
+pub(crate) struct FunctionReturnRequiresX86OrX8664;
+
+#[derive(Diagnostic)]
+#[diag(session_function_return_thunk_extern_requires_non_large_code_model)]
+pub(crate) struct FunctionReturnThunkExternRequiresNonLargeCodeModel;
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
index 4e669c81bf3..68a39de73b7 100644
--- a/compiler/rustc_session/src/options.rs
+++ b/compiler/rustc_session/src/options.rs
@@ -430,6 +430,7 @@ mod desc {
     pub const parse_inlining_threshold: &str =
         "either a boolean (`yes`, `no`, `on`, `off`, etc), or a non-negative number";
     pub const parse_llvm_module_flag: &str = "<key>:<type>:<value>:<behavior>. Type must currently be `u32`. Behavior should be one of (`error`, `warning`, `require`, `override`, `append`, `appendunique`, `max`, `min`)";
+    pub const parse_function_return: &str = "`keep` or `thunk-extern`";
 }
 
 mod parse {
@@ -1359,6 +1360,15 @@ mod parse {
         slot.push((key.to_string(), value, behavior));
         true
     }
+
+    pub(crate) fn parse_function_return(slot: &mut FunctionReturn, v: Option<&str>) -> bool {
+        match v {
+            Some("keep") => *slot = FunctionReturn::Keep,
+            Some("thunk-extern") => *slot = FunctionReturn::ThunkExtern,
+            _ => return false,
+        }
+        true
+    }
 }
 
 options! {
@@ -1603,6 +1613,8 @@ options! {
         "force all crates to be `rustc_private` unstable (default: no)"),
     fuel: Option<(String, u64)> = (None, parse_optimization_fuel, [TRACKED],
         "set the optimization fuel quota for a crate"),
+    function_return: FunctionReturn = (FunctionReturn::default(), parse_function_return, [TRACKED],
+        "replace returns with jumps to `__x86_return_thunk` (default: `keep`)"),
     function_sections: Option<bool> = (None, parse_opt_bool, [TRACKED],
         "whether each function should go in its own section"),
     future_incompat_test: bool = (false, parse_bool, [UNTRACKED],
diff --git a/compiler/rustc_session/src/session.rs b/compiler/rustc_session/src/session.rs
index 20a67d6d036..e79d7648bef 100644
--- a/compiler/rustc_session/src/session.rs
+++ b/compiler/rustc_session/src/session.rs
@@ -1,7 +1,7 @@
 use crate::code_stats::CodeStats;
 pub use crate::code_stats::{DataTypeKind, FieldInfo, FieldKind, SizeKind, VariantInfo};
 use crate::config::{
-    self, CrateType, InstrumentCoverage, OptLevel, OutFileName, OutputType,
+    self, CrateType, FunctionReturn, InstrumentCoverage, OptLevel, OutFileName, OutputType,
     RemapPathScopeComponents, SwitchWithOptPath,
 };
 use crate::config::{ErrorOutputType, Input};
@@ -1678,6 +1678,28 @@ fn validate_commandline_args_with_session_available(sess: &Session) {
             sess.emit_err(errors::IncompatibleLinkerFlavor { flavor, compatible_list });
         }
     }
+
+    if sess.opts.unstable_opts.function_return != FunctionReturn::default() {
+        if sess.target.arch != "x86" && sess.target.arch != "x86_64" {
+            sess.emit_err(errors::FunctionReturnRequiresX86OrX8664);
+        }
+    }
+
+    // The code model check applies to `thunk` and `thunk-extern`, but not `thunk-inline`, so it is
+    // kept as a `match` to force a change if new ones are added, even if we currently only support
+    // `thunk-extern` like Clang.
+    match sess.opts.unstable_opts.function_return {
+        FunctionReturn::Keep => (),
+        FunctionReturn::ThunkExtern => {
+            // FIXME: In principle, the inherited base LLVM target code model could be large,
+            // but this only checks whether we were passed one explicitly (like Clang does).
+            if let Some(code_model) = sess.code_model()
+                && code_model == CodeModel::Large
+            {
+                sess.emit_err(errors::FunctionReturnThunkExternRequiresNonLargeCodeModel);
+            }
+        }
+    }
 }
 
 /// Holds data on the current incremental compilation session, if there is one.