Implement immediate-abort by hooking #[rustc_panic_entrypoint] and lint for functions that look like a panic entrypoint

compiler/rustc_codegen_cranelift/src/base.rs

@@ -14,6 +14,7 @@ use rustc_middle::ty::layout::FnAbiOf;
 use rustc_middle::ty::print::with_no_trimmed_paths;
 use rustc_session::config::OutputFilenames;
 use rustc_span::Symbol;
+use rustc_target::spec::PanicStrategy;
 
 use crate::constant::ConstantCx;
 use crate::debuginfo::{FunctionDebugContext, TypeDebugContext};
@@ -384,6 +385,11 @@ fn codegen_fn_body(fx: &mut FunctionCx<'_, '_, '_>, start_block: Block) {
                 fx.bcx.switch_to_block(failure);
                 fx.bcx.ins().nop();
 
+                if fx.tcx.sess.panic_strategy() == PanicStrategy::ImmediateAbort {
+                    fx.bcx.ins().trap(TrapCode::user(1 /* unreachable */).unwrap());
+                    continue;
+                }
+
                 match &**msg {
                     AssertKind::BoundsCheck { len, index } => {
                         let len = codegen_operand(fx, len).load_scalar(fx);
@@ -1052,6 +1058,10 @@ pub(crate) fn codegen_panic_nounwind<'tcx>(
     msg_str: &str,
     span: Span,
 ) {
+    if fx.tcx.sess.panic_strategy() == PanicStrategy::ImmediateAbort {
+        fx.bcx.ins().trap(TrapCode::user(1 /* unreachable */).unwrap());
+    }
+
     let msg_ptr = crate::constant::pointer_for_anonymous_str(fx, msg_str);
     let msg_len = fx.bcx.ins().iconst(fx.pointer_type, i64::try_from(msg_str.len()).unwrap());
     let args = [msg_ptr, msg_len];

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -4,7 +4,7 @@ use rustc_abi::{Align, ExternAbi};
 use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode};
 use rustc_ast::{LitKind, MetaItem, MetaItemInner};
 use rustc_hir::attrs::{
-    AttributeKind, EiiImplResolution, InlineAttr, Linkage, RtsanSetting, UsedBy,
+    AttributeKind, EiiImplResolution, InlineAttr, Linkage, OptimizeAttr, RtsanSetting, UsedBy,
 };
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LOCAL_CRATE, LocalDefId};
@@ -19,7 +19,7 @@ use rustc_middle::ty::{self as ty, TyCtxt};
 use rustc_session::lint;
 use rustc_session::parse::feature_err;
 use rustc_span::{Span, sym};
-use rustc_target::spec::Os;
+use rustc_target::spec::{Os, PanicStrategy};
 
 use crate::errors;
 use crate::target_features::{
@@ -391,6 +391,24 @@ fn apply_overrides(tcx: TyCtxt<'_>, did: LocalDefId, codegen_fn_attrs: &mut Code
             codegen_fn_attrs.flags |= CodegenFnAttrFlags::NO_MANGLE;
         }
     }
+
+    if tcx.is_panic_entrypoint(did) {
+        // Panic entrypoints are always cold.
+        //
+        // If we have immediate-abort enabled, we want them to be inlined.
+        // They shouldn't be called, but on the off-chance that they are, they should be inlined.
+        //
+        // When the panic strategies that support panic messages are enabled, we want panic
+        // entrypoints outlined and optimized for size.
+        // Most panic entrypoints want #[track_caller] but not all, so we do not add it.
+        codegen_fn_attrs.flags |= CodegenFnAttrFlags::COLD;
+        if tcx.sess.panic_strategy() == PanicStrategy::ImmediateAbort {
+            codegen_fn_attrs.inline = InlineAttr::Always;
+        } else {
+            codegen_fn_attrs.inline = InlineAttr::Never;
+            codegen_fn_attrs.optimize = OptimizeAttr::Size;
+        }
+    }
 }
 
 fn check_result(

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;
@@ -737,6 +738,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

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

@@ -614,6 +614,8 @@ lint_missing_gpu_kernel_export_name = function with the "gpu-kernel" ABI has a m
     .note = mangled names make it hard to find the kernel, this is usually not intended
     .help = use `unsafe(no_mangle)` or `unsafe(export_name = "<name>")`
 
+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/internal.rs

@@ -1,12 +1,14 @@
 //! Some lints that are only useful in the compiler or crates that use compiler internals, such as
 //! Clippy.
 
-use rustc_hir::attrs::AttributeKind;
+use rustc_hir::attrs::{AttributeKind, InlineAttr};
 use rustc_hir::def::Res;
 use rustc_hir::def_id::DefId;
-use rustc_hir::{Expr, ExprKind, HirId, find_attr};
+use rustc_hir::intravisit::FnKind;
+use rustc_hir::{AttrPath, Body, Expr, ExprKind, FnDecl, HirId, find_attr};
 use rustc_middle::ty::{self, GenericArgsRef, PredicatePolarity};
 use rustc_session::{declare_lint_pass, declare_tool_lint};
+use rustc_span::def_id::LocalDefId;
 use rustc_span::hygiene::{ExpnKind, MacroKind};
 use rustc_span::{Span, sym};
 use {rustc_ast as ast, rustc_hir as hir};
@@ -627,3 +629,84 @@ impl EarlyLintPass for ImplicitSysrootCrateImport {
         }
     }
 }
+
+declare_tool_lint! {
+    /// The `missing_panic_entrypoint` lint detects forgotten use of #[rustc_panic_entrypoint].
+    ///
+    /// 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 rustc::MISSING_PANIC_ENTRYPOINT,
+    Allow,
+    "detects missing #[rustc_panic_entrypoint]",
+    report_in_external_macro: true
+}
+
+declare_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] == 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] == sym::rustc_intrinsic {
+                return true;
+            }
+        }
+        false
+    })
+}
+
+impl<'tcx> LateLintPass<'tcx> for MissingPanicEntrypoint {
+    fn check_fn(
+        &mut self,
+        cx: &LateContext<'tcx>,
+        _: FnKind<'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

@@ -240,6 +240,7 @@ late_lint_methods!(
             MapUnitFn: MapUnitFn,
             MissingDebugImplementations: MissingDebugImplementations,
             MissingDoc: MissingDoc,
+            MissingPanicEntrypoint: MissingPanicEntrypoint,
             AsyncClosureUsage: AsyncClosureUsage,
             AsyncFnInTrait: AsyncFnInTrait,
             NonLocalDefinitions: NonLocalDefinitions::default(),
@@ -664,6 +665,7 @@ fn register_internals(store: &mut LintStore) {
     store.register_late_mod_pass(|_| Box::new(SymbolInternStringLiteral));
     store.register_lints(&ImplicitSysrootCrateImport::lint_vec());
     store.register_early_pass(|| Box::new(ImplicitSysrootCrateImport));
+    store.register_late_mod_pass(|_| Box::new(MissingPanicEntrypoint));
     store.register_group(
         false,
         "rustc::internal",
@@ -683,6 +685,7 @@ fn register_internals(store: &mut LintStore) {
             LintId::of(SPAN_USE_EQ_CTXT),
             LintId::of(DIRECT_USE_OF_RUSTC_TYPE_IR),
             LintId::of(IMPLICIT_SYSROOT_CRATE_IMPORT),
+            LintId::of(MISSING_PANIC_ENTRYPOINT),
         ],
     );
 }

compiler/rustc_lint/src/lints.rs

@@ -3325,3 +3325,7 @@ pub(crate) struct UnknownCrateTypesSuggestion {
     pub span: Span,
     pub snippet: Symbol,
 }
+
+#[derive(LintDiagnostic)]
+#[diag(lint_missing_panic_entrypoint)]
+pub(crate) struct MissingPanicEntrypoint;

compiler/rustc_middle/src/ty/context.rs

@@ -49,7 +49,7 @@ use rustc_session::config::CrateType;
 use rustc_session::cstore::{CrateStoreDyn, Untracked};
 use rustc_session::lint::Lint;
 use rustc_span::def_id::{CRATE_DEF_ID, DefPathHash, StableCrateId};
-use rustc_span::{DUMMY_SP, Ident, Span, Symbol, kw};
+use rustc_span::{DUMMY_SP, Ident, Span, Symbol, kw, sym};
 use rustc_type_ir::TyKind::*;
 use rustc_type_ir::lang_items::{SolverAdtLangItem, SolverLangItem, SolverTraitLangItem};
 pub use rustc_type_ir::lift::Lift;
@@ -3556,6 +3556,10 @@ impl<'tcx> TyCtxt<'tcx> {
         }
         false
     }
+
+    pub fn is_panic_entrypoint(self, def_id: impl Into<DefId>) -> bool {
+        self.has_attr(def_id, sym::rustc_panic_entrypoint)
+    }
 }
 
 pub fn provide(providers: &mut Providers) {

compiler/rustc_mir_transform/src/lib.rs

@@ -167,6 +167,7 @@ declare_passes! {
     mod remove_unneeded_drops : RemoveUnneededDrops;
     mod remove_zsts : RemoveZsts;
     mod required_consts : RequiredConstsVisitor;
+    mod panic_entrypoints : PanicEntrypoints;
     mod post_analysis_normalize : PostAnalysisNormalize;
     mod sanity_check : SanityCheck;
     // This pass is public to allow external drivers to perform MIR cleanup
@@ -698,7 +699,7 @@ pub(crate) fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'
             &check_null::CheckNull,
             &check_enums::CheckEnums,
             // Before inlining: trim down MIR with passes to reduce inlining work.
-
+            &panic_entrypoints::PanicEntrypoints,
             // Has to be done before inlining, otherwise actual call will be almost always inlined.
             // Also simple, so can just do first.
             &lower_slice_len::LowerSliceLenCalls,

compiler/rustc_mir_transform/src/panic_entrypoints.rs

@@ -0,0 +1,68 @@
+use rustc_hir::LangItem;
+use rustc_middle::mir::{Const, ConstValue, UnwindAction, *};
+use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_span::Span;
+use rustc_target::spec::PanicStrategy;
+use tracing::{debug, instrument};
+
+pub(super) struct PanicEntrypoints;
+
+impl<'tcx> crate::MirPass<'tcx> for PanicEntrypoints {
+    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
+        sess.panic_strategy() == PanicStrategy::ImmediateAbort
+    }
+
+    #[instrument(level = "trace", skip(self, tcx, body))]
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        if tcx.is_panic_entrypoint(body.source.def_id()) {
+            debug!("Replacing body of {:?}", body.source);
+            let blocks = body.basic_blocks.as_mut();
+            blocks.raw.clear();
+            blocks.push(BasicBlockData::new(Some(abort_terminator(tcx, body.span)), false));
+            return;
+        }
+
+        for bb in body.basic_blocks.as_mut().iter_mut() {
+            let terminator = bb.terminator.as_mut().expect("invalid terminator");
+            let TerminatorKind::Call { func, .. } = &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.is_panic_entrypoint(def_id) {
+                debug!("Remapping call from {:?} to {:?}", body.source, def_id);
+                *terminator = abort_terminator(tcx, terminator.source_info.span);
+            }
+        }
+    }
+
+    fn is_required(&self) -> bool {
+        true
+    }
+}
+
+fn abort_terminator<'tcx>(tcx: TyCtxt<'tcx>, span: Span) -> Terminator<'tcx> {
+    let abort_intrin = tcx.require_lang_item(LangItem::AbortIntrinsic, span);
+    let no_args: [ty::GenericArg<'_>; 0] = [];
+    let func = Operand::Constant(Box::new(ConstOperand {
+        span,
+        user_ty: None,
+        const_: Const::Val(ConstValue::ZeroSized, Ty::new_fn_def(tcx, abort_intrin, no_args)),
+    }));
+    Terminator {
+        source_info: SourceInfo::outermost(span),
+        kind: TerminatorKind::Call {
+            func,
+            args: Box::new([]),
+            destination: Place::return_place(),
+            target: None,
+            unwind: UnwindAction::Unreachable,
+            call_source: CallSource::Misc,
+            fn_span: span,
+        },
+    }
+}

compiler/rustc_passes/src/check_attr.rs

@@ -400,6 +400,7 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                             | sym::rustc_partition_reused
                             | sym::rustc_partition_codegened
                             | sym::rustc_expected_cgu_reuse
+                            | sym::rustc_panic_entrypoint
                             // crate-level attrs, are checked below
                             | sym::feature
                             | sym::register_tool

compiler/rustc_span/src/symbol.rs

@@ -434,6 +434,7 @@ symbols! {
         abi_vectorcall,
         abi_x86_interrupt,
         abort,
+        abort_intrinsic,
         add,
         add_assign,
         add_with_overflow,
@@ -1999,6 +2000,7 @@ symbols! {
         rustc_offload_kernel,
         rustc_on_unimplemented,
         rustc_outlives,
+        rustc_panic_entrypoint,
         rustc_paren_sugar,
         rustc_partition_codegened,
         rustc_partition_reused,