Add rustc_panic_entrypoint

Author: saethlin

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -15,6 +15,7 @@ use rustc_session::config::OptLevel;
 use rustc_span::Span;
 use rustc_span::source_map::Spanned;
 use rustc_target::callconv::{ArgAbi, ArgAttributes, CastTarget, FnAbi, PassMode};
+use rustc_target::spec::PanicStrategy;
 use tracing::{debug, info};
 
 use super::operand::OperandRef;
@@ -730,6 +731,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         bx.switch_to_block(panic_block);
         self.set_debug_loc(bx, terminator.source_info);
 
+        if bx.tcx().sess.panic_strategy() == PanicStrategy::ImmediateAbort {
+            bx.abort();
+            bx.unreachable();
+            return MergingSucc::False;
+        }
+
         // Get the location information.
         let location = self.get_caller_location(bx, terminator.source_info).immediate();
 

compiler/rustc_feature/src/builtin_attrs.rs

@@ -1335,6 +1335,10 @@ pub static BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         "`#[rustc_has_incoherent_inherent_impls]` allows the addition of incoherent inherent impls for \
          the given type by annotating all impl items with `#[rustc_allow_incoherent_impl]`."
     ),
+    rustc_attr!(
+        rustc_panic_entrypoint, Normal, template!(Word), WarnFollowing,
+        EncodeCrossCrate::Yes, "`#[rustc_panic_entrypoint]` makes this function patchable by panic=immediate-abort",
+    ),
 
     BuiltinAttribute {
         name: sym::rustc_diagnostic_item,

compiler/rustc_hir/src/lang_items.rs

@@ -247,6 +247,7 @@ language_item_table! {
     FnMut,                   sym::fn_mut,              fn_mut_trait,               Target::Trait,          GenericRequirement::Exact(1);
     FnOnce,                  sym::fn_once,             fn_once_trait,              Target::Trait,          GenericRequirement::Exact(1);
 
+    AbortIntrinsic,          sym::abort_intrinsic,     abort_intrinsic,            Target::Fn,             GenericRequirement::Exact(0);
     AsyncFn,                 sym::async_fn,            async_fn_trait,             Target::Trait,          GenericRequirement::Exact(1);
     AsyncFnMut,              sym::async_fn_mut,        async_fn_mut_trait,         Target::Trait,          GenericRequirement::Exact(1);
     AsyncFnOnce,             sym::async_fn_once,       async_fn_once_trait,        Target::Trait,          GenericRequirement::Exact(1);

compiler/rustc_lint/messages.ftl

@@ -599,6 +599,8 @@ lint_mismatched_lifetime_syntaxes_suggestion_mixed =
 lint_mismatched_lifetime_syntaxes_suggestion_mixed_only_paths =
     use `'_` for type paths
 
+lint_missing_panic_entrypoint = "diverging functions should usually have #[inline] or #[rustc_panic_entrypoint]"
+
 lint_mixed_script_confusables =
     the usage of Script Group `{$set}` in this crate consists solely of mixed script confusables
     .includes_note = the usage includes {$includes}

compiler/rustc_lint/src/builtin.rs

@@ -25,11 +25,11 @@ use rustc_attr_parsing::AttributeParser;
 use rustc_errors::{Applicability, LintDiagnostic};
 use rustc_feature::GateIssue;
 use rustc_hir as hir;
-use rustc_hir::attrs::{AttributeKind, DocAttribute};
+use rustc_hir::attrs::{AttributeKind, DocAttribute, InlineAttr};
 use rustc_hir::def::{DefKind, Res};
 use rustc_hir::def_id::{CRATE_DEF_ID, DefId, LocalDefId};
 use rustc_hir::intravisit::FnKind as HirFnKind;
-use rustc_hir::{Body, FnDecl, ImplItemImplKind, PatKind, PredicateOrigin, find_attr};
+use rustc_hir::{AttrPath, Body, FnDecl, ImplItemImplKind, PatKind, PredicateOrigin, find_attr};
 use rustc_middle::bug;
 use rustc_middle::lint::LevelAndSource;
 use rustc_middle::ty::layout::LayoutOf;
@@ -3187,3 +3187,97 @@ impl EarlyLintPass for SpecialModuleName {
         }
     }
 }
+
+declare_lint! {
+    /// The `missing_panic_entrypoint` lint detects forgotten use of #[rustc_panic_entrypoint]
+    ///
+    /// ### Example
+    ///
+    /// ```rust,compile_fail
+    /// #[deny(missing_panic_entrypoint)]
+    /// pub fn panic() -> ! {}
+    /// ```
+    ///
+    /// {{produces}}
+    ///
+    /// ### Explanation
+    ///
+    /// This lint is intended to ensure that panic=immediate-abort can function as designed,
+    /// because it uses #[rustc_panic_entrypoint] to locate functions that should be outlined
+    /// for other panic modes, and be deleted entirely when immediate-abort is enabled.
+    pub MISSING_PANIC_ENTRYPOINT,
+    Allow,
+    "detects missing #[rustc_panic_entrypoint]",
+}
+
+#[derive(Default)]
+pub struct MissingPanicEntrypoint;
+
+impl_lint_pass!(MissingPanicEntrypoint => [MISSING_PANIC_ENTRYPOINT]);
+
+fn has_panic_entrypoint(attrs: &[hir::Attribute]) -> bool {
+    attrs.iter().any(|attr| {
+        if let hir::Attribute::Unparsed(box hir::AttrItem {
+            path: AttrPath { segments, .. }, ..
+        }) = attr
+        {
+            if segments[0].name == sym::rustc_panic_entrypoint {
+                return true;
+            }
+        }
+        false
+    })
+}
+
+fn has_inline_encouragement(attrs: &[hir::Attribute]) -> bool {
+    attrs.iter().any(|attr| {
+        matches!(
+            attr,
+            hir::Attribute::Parsed(hir::attrs::AttributeKind::Inline(
+                InlineAttr::Hint | InlineAttr::Always | InlineAttr::Force { .. },
+                _
+            ))
+        )
+    })
+}
+
+fn has_rustc_intrinsic(attrs: &[hir::Attribute]) -> bool {
+    attrs.iter().any(|attr| {
+        if let hir::Attribute::Unparsed(box hir::AttrItem {
+            path: AttrPath { segments, .. }, ..
+        }) = attr
+        {
+            if segments[0].name == sym::rustc_intrinsic {
+                return true;
+            }
+        }
+        false
+    })
+}
+
+impl<'tcx> LateLintPass<'tcx> for MissingPanicEntrypoint {
+    fn check_fn(
+        &mut self,
+        cx: &LateContext<'tcx>,
+        _: HirFnKind<'tcx>,
+        fn_decl: &'tcx FnDecl<'tcx>,
+        _: &'tcx Body<'tcx>,
+        span: Span,
+        def_id: LocalDefId,
+    ) {
+        if matches!(fn_decl.output, hir::FnRetTy::Return(hir::Ty { kind: hir::TyKind::Never, .. }))
+        {
+            let attrs = cx.tcx.hir_attrs(cx.tcx.local_def_id_to_hir_id(def_id));
+            if has_rustc_intrinsic(attrs) {
+                return;
+            }
+            if !has_inline_encouragement(attrs) && !has_panic_entrypoint(attrs) {
+                cx.emit_span_lint(
+                    MISSING_PANIC_ENTRYPOINT,
+                    span,
+                    crate::lints::MissingPanicEntrypoint,
+                );
+            }
+        }
+    }
+}

compiler/rustc_lint/src/lib.rs

@@ -238,6 +238,7 @@ late_lint_methods!(
             MapUnitFn: MapUnitFn,
             MissingDebugImplementations: MissingDebugImplementations,
             MissingDoc: MissingDoc,
+            MissingPanicEntrypoint: MissingPanicEntrypoint,
             AsyncClosureUsage: AsyncClosureUsage,
             AsyncFnInTrait: AsyncFnInTrait,
             NonLocalDefinitions: NonLocalDefinitions::default(),

compiler/rustc_lint/src/lints.rs

@@ -3299,3 +3299,7 @@ pub(crate) struct DocTestUnknown {
 #[derive(LintDiagnostic)]
 #[diag(lint_doc_test_literal)]
 pub(crate) struct DocTestLiteral;
+
+#[derive(LintDiagnostic)]
+#[diag(lint_missing_panic_entrypoint)]
+pub(crate) struct MissingPanicEntrypoint;

compiler/rustc_mir_transform/src/lib.rs

@@ -508,6 +508,52 @@ fn inner_mir_for_ctfe(tcx: TyCtxt<'_>, def: LocalDefId) -> Body<'_> {
     body
 }
 
+fn remap_mir_for_panic_immediate_abort<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    mut body: Body<'tcx>,
+) -> Body<'tcx> {
+    use rustc_middle::mir::{Const, ConstValue, UnwindAction};
+    use rustc_middle::ty::Ty;
+    use rustc_target::spec::PanicStrategy;
+
+    if tcx.sess.panic_strategy() != PanicStrategy::ImmediateAbort {
+        return body;
+    }
+
+    let make_abort_intrinsic_operand = || {
+        let abort_intrin = tcx.lang_items().abort_intrinsic().unwrap();
+        let no_args: [ty::GenericArg<'_>; 0] = [];
+        Operand::Constant(Box::new(ConstOperand {
+            span: DUMMY_SP,
+            user_ty: None,
+            const_: Const::Val(ConstValue::ZeroSized, Ty::new_fn_def(tcx, abort_intrin, no_args)),
+        }))
+    };
+
+    for bb in body.basic_blocks.as_mut().iter_mut() {
+        let terminator = bb.terminator.as_mut().expect("invalid terminator");
+        let TerminatorKind::Call { func, args, unwind, .. } = &mut terminator.kind else {
+            continue;
+        };
+        let Operand::Constant(box ConstOperand { const_, .. }) = &func else {
+            continue;
+        };
+        let ty::FnDef(def_id, _) = *const_.ty().kind() else {
+            continue;
+        };
+        if !tcx.has_attr(def_id, sym::rustc_panic_entrypoint) {
+            continue;
+        }
+
+        debug!("Remapping call from {:?} to {:?}", body.source, def_id);
+
+        *func = make_abort_intrinsic_operand();
+        *args = Box::new([]);
+        *unwind = UnwindAction::Unreachable;
+    }
+    body
+}
+
 /// Obtain just the main MIR (no promoteds) and run some cleanups on it. This also runs
 /// mir borrowck *before* doing so in order to ensure that borrowck can be run and doesn't
 /// end up missing the source MIR due to stealing happening.
@@ -796,7 +842,8 @@ fn inner_optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> Body<'_> {
     }
     debug!("about to call mir_drops_elaborated...");
     let body = tcx.mir_drops_elaborated_and_const_checked(did).steal();
-    let mut body = remap_mir_for_const_eval_select(tcx, body, hir::Constness::NotConst);
+    let body = remap_mir_for_const_eval_select(tcx, body, hir::Constness::NotConst);
+    let mut body = remap_mir_for_panic_immediate_abort(tcx, body);
 
     if body.tainted_by_errors.is_some() {
         return body;

compiler/rustc_span/src/symbol.rs

@@ -424,6 +424,7 @@ symbols! {
         abi_vectorcall,
         abi_x86_interrupt,
         abort,
+        abort_intrinsic,
         add,
         add_assign,
         add_with_overflow,
@@ -1981,6 +1982,7 @@ symbols! {
         rustc_offload_kernel,
         rustc_on_unimplemented,
         rustc_outlives,
+        rustc_panic_entrypoint,
         rustc_paren_sugar,
         rustc_partition_codegened,
         rustc_partition_reused,

library/alloc/src/alloc.rs

@@ -396,7 +396,9 @@ unsafe extern "Rust" {
 #[cfg(not(no_global_oom_handling))]
 #[cold]
 #[optimize(size)]
+#[rustc_panic_entrypoint]
 pub const fn handle_alloc_error(layout: Layout) -> ! {
+    #[inline]
     const fn ct_error(_: Layout) -> ! {
         panic!("allocation failed");
     }
@@ -408,13 +410,7 @@ pub const fn handle_alloc_error(layout: Layout) -> ! {
         }
     }
 
-    #[cfg(not(panic = "immediate-abort"))]
-    {
-        core::intrinsics::const_eval_select((layout,), ct_error, rt_error)
-    }
-
-    #[cfg(panic = "immediate-abort")]
-    ct_error(layout)
+    core::intrinsics::const_eval_select((layout,), ct_error, rt_error)
 }
 
 #[cfg(not(no_global_oom_handling))]
@@ -425,6 +421,7 @@ pub mod __alloc_error_handler {
     // called via generated `__rust_alloc_error_handler` if there is no
     // `#[alloc_error_handler]`.
     #[rustc_std_internal_symbol]
+    #[allow(missing_panic_entrypoint)]
     pub unsafe fn __rdl_alloc_error_handler(size: usize, _align: usize) -> ! {
         core::panicking::panic_nounwind_fmt(
             format_args!("memory allocation of {size} bytes failed"),