Auto merge of #134082 - davidtwco:forced-inlining, r=saethlin

mir_transform: implement `#[rustc_force_inline]`

Adds `#[rustc_force_inline]` which is similar to always inlining but reports an error if the inlining was not possible.

- `#[rustc_force_inline]` can only be applied to free functions to guarantee that the MIR inliner will be able to resolve calls.
- `rustc_mir_transform::inline::Inline` is refactored into two passes (`Inline` and `ForceInline`), sharing the vast majority of the implementation.
  - `rustc_mir_transform::inline::ForceInline` can't be disabled so annotated items are always inlined.
  - `rustc_mir_transform::inline::ForceInline` runs regardless of optimisation level.
- `#[rustc_force_inline]` won't inline unless target features match, as with normal inlining.
- MIR validation will ICE if a `#[rustc_force_inline]` isn't inlined, to guarantee that it will never be codegened independently. As a further guarantee, monomorphisation collection will always decide that `#[rustc_force_inline]` functions cannot be codegened locally.
- Like other intrinsics, `#[rustc_force_inline]` annotated functions cannot be cast to function pointers.
- As with other rustc attrs, this cannot be used by users, just within the compiler and standard library.
- This is only implemented within rustc, so should avoid any limitations of LLVM's inlining.

It is intended that this attribute be used with intrinsics that must be inlined for security reasons. For example, pointer authentication intrinsics would allow Rust users to make use of pointer authentication instructions, but if these intrinsic functions were in the binary then they could be used as gadgets with ROP attacks, defeating the point of introducing them. We don't have any intrinsics like this today, but I expect to upstream some once a force inlining mechanism such as this is available.

cc #131687 rust-lang/rfcs#3711 - this approach should resolve the concerns from these previous attempts

r? `@saethlin`

Cargo.lock

@@ -4158,6 +4158,7 @@ dependencies = [
  "rustc_apfloat",
  "rustc_arena",
  "rustc_ast",
+ "rustc_attr_parsing",
  "rustc_data_structures",
  "rustc_errors",
  "rustc_fluent_macro",

compiler/rustc_attr_data_structures/src/attributes.rs

@@ -11,6 +11,22 @@ pub enum InlineAttr {
     Hint,
     Always,
     Never,
+    /// `#[rustc_force_inline]` forces inlining to happen in the MIR inliner - it reports an error
+    /// if the inlining cannot happen. It is limited to only free functions so that the calls
+    /// can always be resolved.
+    Force {
+        attr_span: Span,
+        reason: Option<Symbol>,
+    },
+}
+
+impl InlineAttr {
+    pub fn always(&self) -> bool {
+        match self {
+            InlineAttr::Always | InlineAttr::Force { .. } => true,
+            InlineAttr::None | InlineAttr::Hint | InlineAttr::Never => false,
+        }
+    }
 }
 
 #[derive(Clone, Encodable, Decodable, Debug, PartialEq, Eq, HashStable_Generic)]

compiler/rustc_codegen_gcc/src/attributes.rs

@@ -20,7 +20,7 @@ fn inline_attr<'gcc, 'tcx>(
 ) -> Option<FnAttribute<'gcc>> {
     match inline {
         InlineAttr::Hint => Some(FnAttribute::Inline),
-        InlineAttr::Always => Some(FnAttribute::AlwaysInline),
+        InlineAttr::Always | InlineAttr::Force { .. } => Some(FnAttribute::AlwaysInline),
         InlineAttr::Never => {
             if cx.sess().target.arch != "amdgpu" {
                 Some(FnAttribute::NoInline)

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -37,7 +37,9 @@ fn inline_attr<'ll>(cx: &CodegenCx<'ll, '_>, inline: InlineAttr) -> Option<&'ll
     }
     match inline {
         InlineAttr::Hint => Some(AttributeKind::InlineHint.create_attr(cx.llcx)),
-        InlineAttr::Always => Some(AttributeKind::AlwaysInline.create_attr(cx.llcx)),
+        InlineAttr::Always | InlineAttr::Force { .. } => {
+            Some(AttributeKind::AlwaysInline.create_attr(cx.llcx))
+        }
         InlineAttr::Never => {
             if cx.sess().target.arch != "amdgpu" {
                 Some(AttributeKind::NoInline.create_attr(cx.llcx))

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -18,6 +18,7 @@ use rustc_session::parse::feature_err;
 use rustc_session::{Session, lint};
 use rustc_span::{Ident, Span, sym};
 use rustc_target::spec::{SanitizerSet, abi};
+use tracing::debug;
 
 use crate::errors;
 use crate::target_features::{check_target_feature_trait_unsafe, from_target_feature_attr};
@@ -525,6 +526,7 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
         if !attr.has_name(sym::inline) {
             return ia;
         }
+
         if attr.is_word() {
             InlineAttr::Hint
         } else if let Some(ref items) = attr.meta_item_list() {
@@ -547,6 +549,20 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
             ia
         }
     });
+    codegen_fn_attrs.inline = attrs.iter().fold(codegen_fn_attrs.inline, |ia, attr| {
+        if !attr.has_name(sym::rustc_force_inline) || !tcx.features().rustc_attrs() {
+            return ia;
+        }
+
+        if attr.is_word() {
+            InlineAttr::Force { attr_span: attr.span, reason: None }
+        } else if let Some(val) = attr.value_str() {
+            InlineAttr::Force { attr_span: attr.span, reason: Some(val) }
+        } else {
+            debug!("`rustc_force_inline` not checked by attribute validation");
+            ia
+        }
+    });
 
     // naked function MUST NOT be inlined! This attribute is required for the rust compiler itself,
     // but not for the code generation backend because at that point the naked function will just be
@@ -596,7 +612,7 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
     // is probably a poor usage of `#[inline(always)]` and easily avoided by not using the attribute.
     if tcx.features().target_feature_11()
         && tcx.is_closure_like(did.to_def_id())
-        && codegen_fn_attrs.inline != InlineAttr::Always
+        && !codegen_fn_attrs.inline.always()
     {
         let owner_id = tcx.parent(did.to_def_id());
         if tcx.def_kind(owner_id).has_codegen_attrs() {
@@ -606,22 +622,28 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
         }
     }
 
-    // If a function uses #[target_feature] it can't be inlined into general
+    // If a function uses `#[target_feature]` it can't be inlined into general
     // purpose functions as they wouldn't have the right target features
-    // enabled. For that reason we also forbid #[inline(always)] as it can't be
+    // enabled. For that reason we also forbid `#[inline(always)]` as it can't be
     // respected.
-    if !codegen_fn_attrs.target_features.is_empty() && codegen_fn_attrs.inline == InlineAttr::Always
+    //
+    // `#[rustc_force_inline]` doesn't need to be prohibited here, only
+    // `#[inline(always)]`, as forced inlining is implemented entirely within
+    // rustc (and so the MIR inliner can do any necessary checks for compatible target
+    // features).
+    //
+    // This sidesteps the LLVM blockers in enabling `target_features` +
+    // `inline(always)` to be used together (see rust-lang/rust#116573 and
+    // llvm/llvm-project#70563).
+    if !codegen_fn_attrs.target_features.is_empty()
+        && matches!(codegen_fn_attrs.inline, InlineAttr::Always)
     {
         if let Some(span) = inline_span {
-            tcx.dcx().span_err(
-                span,
-                "cannot use `#[inline(always)]` with \
-                     `#[target_feature]`",
-            );
+            tcx.dcx().span_err(span, "cannot use `#[inline(always)]` with `#[target_feature]`");
         }
     }
 
-    if !codegen_fn_attrs.no_sanitize.is_empty() && codegen_fn_attrs.inline == InlineAttr::Always {
+    if !codegen_fn_attrs.no_sanitize.is_empty() && codegen_fn_attrs.inline.always() {
         if let (Some(no_sanitize_span), Some(inline_span)) = (no_sanitize_span, inline_span) {
             let hir_id = tcx.local_def_id_to_hir_id(did);
             tcx.node_span_lint(

compiler/rustc_feature/src/builtin_attrs.rs

@@ -1019,6 +1019,10 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         rustc_no_mir_inline, Normal, template!(Word), WarnFollowing, EncodeCrossCrate::Yes,
         "#[rustc_no_mir_inline] prevents the MIR inliner from inlining a function while not affecting codegen"
     ),
+    rustc_attr!(
+        rustc_force_inline, Normal, template!(Word, NameValueStr: "reason"), WarnFollowing, EncodeCrossCrate::Yes,
+        "#![rustc_force_inline] forces a free function to be inlined"
+    ),
 
     // ==========================================================================
     // Internal attributes, Testing:

compiler/rustc_hir_typeck/src/coercion.rs

@@ -38,6 +38,7 @@
 use std::ops::Deref;
 
 use rustc_abi::ExternAbi;
+use rustc_attr_parsing::InlineAttr;
 use rustc_errors::codes::*;
 use rustc_errors::{Applicability, Diag, struct_span_code_err};
 use rustc_hir as hir;
@@ -926,8 +927,13 @@ impl<'f, 'tcx> Coerce<'f, 'tcx> {
                         return Err(TypeError::IntrinsicCast);
                     }
 
-                    // Safe `#[target_feature]` functions are not assignable to safe fn pointers (RFC 2396).
+                    let fn_attrs = self.tcx.codegen_fn_attrs(def_id);
+                    if matches!(fn_attrs.inline, InlineAttr::Force { .. }) {
+                        return Err(TypeError::ForceInlineCast);
+                    }
 
+                    // Safe `#[target_feature]` functions are not assignable to safe fn pointers
+                    // (RFC 2396).
                     if b_hdr.safety.is_safe()
                         && !self.tcx.codegen_fn_attrs(def_id).target_features.is_empty()
                     {
@@ -1197,6 +1203,17 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             return Ok(prev_ty);
         }
 
+        let is_force_inline = |ty: Ty<'tcx>| {
+            if let ty::FnDef(did, _) = ty.kind() {
+                matches!(self.tcx.codegen_fn_attrs(did).inline, InlineAttr::Force { .. })
+            } else {
+                false
+            }
+        };
+        if is_force_inline(prev_ty) || is_force_inline(new_ty) {
+            return Err(TypeError::ForceInlineCast);
+        }
+
         // Special-case that coercion alone cannot handle:
         // Function items or non-capturing closures of differing IDs or GenericArgs.
         let (a_sig, b_sig) = {

compiler/rustc_middle/src/mir/mono.rs

@@ -132,9 +132,10 @@ impl<'tcx> MonoItem<'tcx> {
                 // creating one copy of this `#[inline]` function which may
                 // conflict with upstream crates as it could be an exported
                 // symbol.
-                match tcx.codegen_fn_attrs(instance.def_id()).inline {
-                    InlineAttr::Always => InstantiationMode::LocalCopy,
-                    _ => InstantiationMode::GloballyShared { may_conflict: true },
+                if tcx.codegen_fn_attrs(instance.def_id()).inline.always() {
+                    InstantiationMode::LocalCopy
+                } else {
+                    InstantiationMode::GloballyShared { may_conflict: true }
                 }
             }
             MonoItem::Static(..) | MonoItem::GlobalAsm(..) => {

compiler/rustc_middle/src/ty/error.rs

@@ -109,6 +109,9 @@ impl<'tcx> TypeError<'tcx> {
             TypeError::ConstMismatch(ref values) => {
                 format!("expected `{}`, found `{}`", values.expected, values.found).into()
             }
+            TypeError::ForceInlineCast => {
+                "cannot coerce functions which must be inlined to function pointers".into()
+            }
             TypeError::IntrinsicCast => "cannot coerce intrinsics to function pointers".into(),
             TypeError::TargetFeatureCast(_) => {
                 "cannot coerce functions with `#[target_feature]` to safe function pointers".into()

compiler/rustc_mir_build/Cargo.toml

@@ -12,6 +12,7 @@ rustc_abi = { path = "../rustc_abi" }
 rustc_apfloat = "0.2.0"
 rustc_arena = { path = "../rustc_arena" }
 rustc_ast = { path = "../rustc_ast" }
+rustc_attr_parsing = { path = "../rustc_attr_parsing" }
 rustc_data_structures = { path = "../rustc_data_structures" }
 rustc_errors = { path = "../rustc_errors" }
 rustc_fluent_macro = { path = "../rustc_fluent_macro" }

compiler/rustc_mir_build/messages.ftl

@@ -118,6 +118,12 @@ mir_build_extern_static_requires_unsafe_unsafe_op_in_unsafe_fn_allowed =
     .note = extern statics are not controlled by the Rust type system: invalid data, aliasing violations or data races will cause undefined behavior
     .label = use of extern static
 
+mir_build_force_inline =
+    `{$callee}` is incompatible with `#[rustc_force_inline]`
+    .attr = annotation here
+    .callee = `{$callee}` defined here
+    .note = incompatible due to: {$reason}
+
 mir_build_inform_irrefutable = `let` bindings require an "irrefutable pattern", like a `struct` or an `enum` with only one variant
 
 mir_build_initializing_type_with_requires_unsafe =

compiler/rustc_mir_build/src/builder/mod.rs

@@ -29,6 +29,7 @@ use rustc_span::{Span, Symbol, sym};
 use super::lints;
 use crate::builder::expr::as_place::PlaceBuilder;
 use crate::builder::scope::DropKind;
+use crate::check_inline;
 
 pub(crate) fn closure_saved_names_of_captured_variables<'tcx>(
     tcx: TyCtxt<'tcx>,
@@ -80,6 +81,7 @@ pub(crate) fn mir_build<'tcx>(tcx: TyCtxtAt<'tcx>, def: LocalDefId) -> Body<'tcx
     };
 
     lints::check(tcx, &body);
+    check_inline::check_force_inline(tcx, &body);
 
     // The borrow checker will replace all the regions here with its own
     // inference variables. There's no point having non-erased regions here.

compiler/rustc_mir_build/src/check_inline.rs

@@ -0,0 +1,81 @@
+use rustc_attr_parsing::InlineAttr;
+use rustc_hir::def_id::DefId;
+use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
+use rustc_middle::mir::{Body, TerminatorKind};
+use rustc_middle::ty;
+use rustc_middle::ty::TyCtxt;
+use rustc_span::sym;
+
+/// Check that a body annotated with `#[rustc_force_inline]` will not fail to inline based on its
+/// definition alone (irrespective of any specific caller).
+pub(crate) fn check_force_inline<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) {
+    let def_id = body.source.def_id();
+    if !tcx.hir().body_owner_kind(def_id).is_fn_or_closure() || !def_id.is_local() {
+        return;
+    }
+    let InlineAttr::Force { attr_span, .. } = tcx.codegen_fn_attrs(def_id).inline else {
+        return;
+    };
+
+    if let Err(reason) =
+        is_inline_valid_on_fn(tcx, def_id).and_then(|_| is_inline_valid_on_body(tcx, body))
+    {
+        tcx.dcx().emit_err(crate::errors::InvalidForceInline {
+            attr_span,
+            callee_span: tcx.def_span(def_id),
+            callee: tcx.def_path_str(def_id),
+            reason,
+        });
+    }
+}
+
+pub fn is_inline_valid_on_fn<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Result<(), &'static str> {
+    let codegen_attrs = tcx.codegen_fn_attrs(def_id);
+    if tcx.has_attr(def_id, sym::rustc_no_mir_inline) {
+        return Err("#[rustc_no_mir_inline]");
+    }
+
+    // FIXME(#127234): Coverage instrumentation currently doesn't handle inlined
+    // MIR correctly when Modified Condition/Decision Coverage is enabled.
+    if tcx.sess.instrument_coverage_mcdc() {
+        return Err("incompatible with MC/DC coverage");
+    }
+
+    let ty = tcx.type_of(def_id);
+    if match ty.instantiate_identity().kind() {
+        ty::FnDef(..) => tcx.fn_sig(def_id).instantiate_identity().c_variadic(),
+        ty::Closure(_, args) => args.as_closure().sig().c_variadic(),
+        _ => false,
+    } {
+        return Err("C variadic");
+    }
+
+    if codegen_attrs.flags.contains(CodegenFnAttrFlags::COLD) {
+        return Err("cold");
+    }
+
+    // Intrinsic fallback bodies are automatically made cross-crate inlineable,
+    // but at this stage we don't know whether codegen knows the intrinsic,
+    // so just conservatively don't inline it. This also ensures that we do not
+    // accidentally inline the body of an intrinsic that *must* be overridden.
+    if tcx.has_attr(def_id, sym::rustc_intrinsic) {
+        return Err("callee is an intrinsic");
+    }
+
+    Ok(())
+}
+
+pub fn is_inline_valid_on_body<'tcx>(
+    _: TyCtxt<'tcx>,
+    body: &Body<'tcx>,
+) -> Result<(), &'static str> {
+    if body
+        .basic_blocks
+        .iter()
+        .any(|bb| matches!(bb.terminator().kind, TerminatorKind::TailCall { .. }))
+    {
+        return Err("can't inline functions with tail calls");
+    }
+
+    Ok(())
+}

compiler/rustc_mir_build/src/errors.rs

@@ -1107,3 +1107,15 @@ impl<'a> Subdiagnostic for Rust2024IncompatiblePatSugg<'a> {
         );
     }
 }
+
+#[derive(Diagnostic)]
+#[diag(mir_build_force_inline)]
+#[note]
+pub(crate) struct InvalidForceInline {
+    #[primary_span]
+    pub attr_span: Span,
+    #[label(mir_build_callee)]
+    pub callee_span: Span,
+    pub callee: String,
+    pub reason: &'static str,
+}

compiler/rustc_mir_build/src/lib.rs

@@ -15,6 +15,7 @@
 // "Go to file" feature to silently ignore all files in the module, probably
 // because it assumes that "build" is a build-output directory. See #134365.
 mod builder;
+pub mod check_inline;
 mod check_tail_calls;
 mod check_unsafety;
 mod errors;

compiler/rustc_mir_transform/messages.ftl

@@ -19,6 +19,17 @@ mir_transform_ffi_unwind_call = call to {$foreign ->
 mir_transform_fn_item_ref = taking a reference to a function item does not give a function pointer
     .suggestion = cast `{$ident}` to obtain a function pointer
 
+mir_transform_force_inline =
+    `{$callee}` could not be inlined into `{$caller}` but is required to be inlined
+    .call = ...`{$callee}` called here
+    .attr = inlining due to this annotation
+    .caller = within `{$caller}`...
+    .callee = `{$callee}` defined here
+    .note = could not be inlined due to: {$reason}
+
+mir_transform_force_inline_justification =
+    `{$callee}` is required to be inlined to: {$sym}
+
 mir_transform_must_not_suspend = {$pre}`{$def_path}`{$post} held across a suspend point, but should not be
     .label = the value is held across this suspend point
     .note = {$reason}

compiler/rustc_mir_transform/src/cross_crate_inline.rs

@@ -46,7 +46,7 @@ fn cross_crate_inlinable(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
     // #[inline(never)] to force code generation.
     match codegen_fn_attrs.inline {
         InlineAttr::Never => return false,
-        InlineAttr::Hint | InlineAttr::Always => return true,
+        InlineAttr::Hint | InlineAttr::Always | InlineAttr::Force { .. } => return true,
         _ => {}
     }
 
@@ -69,8 +69,9 @@ fn cross_crate_inlinable(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
     // Don't do any inference if codegen optimizations are disabled and also MIR inlining is not
     // enabled. This ensures that we do inference even if someone only passes -Zinline-mir,
     // which is less confusing than having to also enable -Copt-level=1.
-    if matches!(tcx.sess.opts.optimize, OptLevel::No) && !pm::should_run_pass(tcx, &inline::Inline)
-    {
+    let inliner_will_run = pm::should_run_pass(tcx, &inline::Inline)
+        || inline::ForceInline::should_run_pass_for_callee(tcx, def_id.to_def_id());
+    if matches!(tcx.sess.opts.optimize, OptLevel::No) && !inliner_will_run {
         return false;
     }