mir_transform: implement forced inlining

Adds `#[rustc_force_inline]` which is similar to always inlining but
reports an error if the inlining was not possible, and which always
attempts to inline annotated items, regardless of optimisation levels.
It can only be applied to free functions to guarantee that the MIR
inliner will be able to resolve calls.

Author: davidtwco

compiler/rustc_attr/src/builtin.rs

@@ -47,6 +47,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

@@ -19,6 +19,7 @@ use rustc_session::{Session, lint};
 use rustc_span::symbol::Ident;
 use rustc_span::{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};
@@ -514,31 +515,50 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
     }
 
     codegen_fn_attrs.inline = attrs.iter().fold(InlineAttr::None, |ia, attr| {
-        if !attr.has_name(sym::inline) {
-            return ia;
-        }
-        match attr.meta_kind() {
-            Some(MetaItemKind::Word) => InlineAttr::Hint,
-            Some(MetaItemKind::List(ref items)) => {
-                inline_span = Some(attr.span);
-                if items.len() != 1 {
-                    struct_span_code_err!(tcx.dcx(), attr.span, E0534, "expected one argument")
-                        .emit();
-                    InlineAttr::None
-                } else if list_contains_name(items, sym::always) {
-                    InlineAttr::Always
-                } else if list_contains_name(items, sym::never) {
-                    InlineAttr::Never
-                } else {
-                    struct_span_code_err!(tcx.dcx(), items[0].span(), E0535, "invalid argument")
+        if attr.has_name(sym::inline) {
+            match attr.meta_kind() {
+                Some(MetaItemKind::Word) => InlineAttr::Hint,
+                Some(MetaItemKind::List(ref items)) => {
+                    inline_span = Some(attr.span);
+                    if items.len() != 1 {
+                        struct_span_code_err!(tcx.dcx(), attr.span, E0534, "expected one argument")
+                            .emit();
+                        InlineAttr::None
+                    } else if list_contains_name(items, sym::always) {
+                        InlineAttr::Always
+                    } else if list_contains_name(items, sym::never) {
+                        InlineAttr::Never
+                    } else {
+                        struct_span_code_err!(
+                            tcx.dcx(),
+                            items[0].span(),
+                            E0535,
+                            "invalid argument"
+                        )
                         .with_help("valid inline arguments are `always` and `never`")
                         .emit();
 
-                    InlineAttr::None
+                        InlineAttr::None
+                    }
                 }
+                Some(MetaItemKind::NameValue(_)) => ia,
+                None => ia,
             }
-            Some(MetaItemKind::NameValue(_)) => ia,
-            None => ia,
+        } else if attr.has_name(sym::rustc_force_inline) && tcx.features().rustc_attrs() {
+            match attr.meta_kind() {
+                Some(MetaItemKind::NameValue(lit)) => {
+                    InlineAttr::Force { attr_span: attr.span, reason: Some(lit.symbol) }
+                }
+                Some(MetaItemKind::Word) => {
+                    InlineAttr::Force { attr_span: attr.span, reason: None }
+                }
+                _ => {
+                    debug!("`rustc_force_inline` not checked by attribute validation");
+                    ia
+                }
+            }
+        } else {
+            ia
         }
     });
 
@@ -586,7 +606,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() {
@@ -596,11 +616,15 @@ 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, that
+    // is implemented entirely in rustc can attempt to inline and error if it cannot.
+    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(
@@ -611,7 +635,7 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
         }
     }
 
-    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

@@ -1014,6 +1014,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::InlineAttr;
 use rustc_errors::codes::*;
 use rustc_errors::{Applicability, Diag, struct_span_code_err};
 use rustc_hir as hir;
@@ -923,11 +924,14 @@ 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, rustc_attr::InlineAttr::Force { .. }) {
+                        return Err(TypeError::ForceInlineCast);
+                    }
 
-                    if b_hdr.safety == hir::Safety::Safe
-                        && !self.tcx.codegen_fn_attrs(def_id).target_features.is_empty()
-                    {
+                    // Safe `#[target_feature]` functions are not assignable to safe fn pointers
+                    // (RFC 2396).
+                    if b_hdr.safety == hir::Safety::Safe && !fn_attrs.target_features.is_empty() {
                         return Err(TypeError::TargetFeatureCast(def_id));
                     }
                 }
@@ -1194,6 +1198,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_lint_defs/src/builtin.rs

@@ -46,6 +46,7 @@ declare_lint_pass! {
         EXPORTED_PRIVATE_DEPENDENCIES,
         FFI_UNWIND_CALLS,
         FORBIDDEN_LINT_GROUPS,
+        FORCED_INLINE,
         FUNCTION_ITEM_REFERENCES,
         FUZZY_PROVENANCE_CASTS,
         HIDDEN_GLOB_REEXPORTS,
@@ -5178,3 +5179,31 @@ declare_lint! {
         reference: "issue #116558 <https://github.com/rust-lang/rust/issues/116558>",
     };
 }
+
+declare_lint! {
+    /// The `forced_inline` lint is emitted when a function annotated with
+    /// `#[rustc_force_inline]` was not able to be inlined.
+    ///
+    /// ### Example
+    ///
+    /// ```rust,ignore (needs rustc_attrs)
+    /// #[rustc_no_mir_inline]
+    /// #[rustc_force_inline]
+    /// fn foo() { }
+    ///
+    /// fn bar() { foo() }
+    /// ```
+    ///
+    /// ### Explanation
+    ///
+    /// Functions can be marked as `#[rustc_force_inline]` in the standard
+    /// library if they are required to be inlined in order to uphold
+    /// security properties or some other similar guarantee.
+    ///
+    /// In some circumstances, these functions cannot be inlined and a
+    /// reason will be provided, this can either be rectified or the
+    /// lint can be silenced if the risk is acceptable.
+    pub FORCED_INLINE,
+    Deny,
+    "`#[rustc_force_inline]`-annotated function could not be inlined"
+}

compiler/rustc_middle/src/mir/mono.rs

@@ -1,7 +1,6 @@
 use std::fmt;
 use std::hash::Hash;
 
-use rustc_attr::InlineAttr;
 use rustc_data_structures::base_n::{BaseNString, CASE_INSENSITIVE, ToBaseN};
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::FxIndexMap;
@@ -111,7 +110,8 @@ impl<'tcx> MonoItem<'tcx> {
                     return InstantiationMode::GloballyShared { may_conflict: false };
                 }
 
-                if let InlineAttr::Never = tcx.codegen_fn_attrs(instance.def_id()).inline
+                if let rustc_attr::InlineAttr::Never =
+                    tcx.codegen_fn_attrs(instance.def_id()).inline
                     && self.is_generic_fn()
                 {
                     // Upgrade inline(never) to a globally shared instance.
@@ -130,9 +130,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_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;
     }
 

compiler/rustc_mir_transform/src/errors.rs

@@ -4,8 +4,8 @@ use rustc_macros::{Diagnostic, LintDiagnostic, Subdiagnostic};
 use rustc_middle::mir::AssertKind;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::lint::{self, Lint};
-use rustc_span::Span;
 use rustc_span::def_id::DefId;
+use rustc_span::{Span, Symbol};
 
 use crate::fluent_generated as fluent;
 
@@ -142,3 +142,29 @@ pub(crate) struct MustNotSuspendReason {
 #[note(mir_transform_note2)]
 #[help]
 pub(crate) struct UndefinedTransmute;
+
+#[derive(Diagnostic)]
+#[diag(mir_transform_force_inline)]
+#[note]
+pub(crate) struct ForceInlineFailure {
+    #[label(mir_transform_caller)]
+    pub caller_span: Span,
+    #[label(mir_transform_callee)]
+    pub callee_span: Span,
+    #[label(mir_transform_attr)]
+    pub attr_span: Span,
+    #[primary_span]
+    #[label(mir_transform_call)]
+    pub call_span: Span,
+    pub callee: String,
+    pub caller: String,
+    pub reason: &'static str,
+    #[subdiagnostic]
+    pub justification: Option<ForceInlineJustification>,
+}
+
+#[derive(Subdiagnostic)]
+#[note(mir_transform_force_inline_justification)]
+pub(crate) struct ForceInlineJustification {
+    pub sym: Symbol,
+}