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

Cargo.lock

@@ -670,7 +670,7 @@ dependencies = [
  "tempfile",
  "termize",
  "toml 0.9.8",
- "ui_test",
+ "ui_test 0.30.3",
  "walkdir",
 ]
 
@@ -892,6 +892,18 @@ dependencies = [
  "windows-sys 0.59.0",
 ]
 
+[[package]]
+name = "console"
+version = "0.16.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d64e8af5551369d19cf50138de61f1c42074ab970f74e99be916646777f8fc87"
+dependencies = [
+ "encode_unicode",
+ "libc",
+ "unicode-width 0.2.2",
+ "windows-sys 0.61.2",
+]
+
 [[package]]
 name = "constant_time_eq"
 version = "0.3.1"
@@ -1969,13 +1981,26 @@ version = "0.17.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235"
 dependencies = [
- "console",
+ "console 0.15.11",
  "number_prefix",
  "portable-atomic",
  "unicode-width 0.2.2",
  "web-time",
 ]
 
+[[package]]
+name = "indicatif"
+version = "0.18.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9375e112e4b463ec1b1c6c011953545c65a30164fbab5b581df32b3abf0dcb88"
+dependencies = [
+ "console 0.16.3",
+ "portable-atomic",
+ "unicode-width 0.2.2",
+ "unit-prefix",
+ "web-time",
+]
+
 [[package]]
 name = "indoc"
 version = "1.0.9"
@@ -2501,7 +2526,7 @@ dependencies = [
  "serde_json",
  "smallvec",
  "tempfile",
- "ui_test",
+ "ui_test 0.30.4",
 ]
 
 [[package]]
@@ -5531,7 +5556,7 @@ checksum = "8f50febec83f5ee1df3015341d8bd429f2d1cc62bcba7ea2076759d315084683"
 name = "test-float-parse"
 version = "0.1.0"
 dependencies = [
- "indicatif",
+ "indicatif 0.17.11",
  "num",
  "rand 0.9.2",
  "rand_chacha 0.9.0",
@@ -5969,7 +5994,32 @@ dependencies = [
  "colored 3.0.0",
  "comma",
  "crossbeam-channel",
- "indicatif",
+ "indicatif 0.17.11",
+ "levenshtein",
+ "prettydiff",
+ "regex",
+ "rustc_version",
+ "rustfix",
+ "serde",
+ "serde_json",
+ "spanned",
+]
+
+[[package]]
+name = "ui_test"
+version = "0.30.4"
+source = "git+https://github.com/saethlin/ui_test?branch=target-in-compile-flags#62e278f072527fa19477aaedcd6b6fcd5ce1a485"
+dependencies = [
+ "annotate-snippets 0.11.5",
+ "anyhow",
+ "bstr",
+ "cargo-platform 0.1.9",
+ "cargo_metadata 0.18.1",
+ "color-eyre",
+ "colored 3.0.0",
+ "comma",
+ "crossbeam-channel",
+ "indicatif 0.18.3",
  "levenshtein",
  "prettydiff",
  "regex",
@@ -6106,6 +6156,12 @@ dependencies = [
  "diff",
 ]
 
+[[package]]
+name = "unit-prefix"
+version = "0.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "81e544489bf3d8ef66c953931f56617f423cd4b5494be343d9b9d3dda037b9a3"
+
 [[package]]
 name = "unstable-book-gen"
 version = "0.1.0"

compiler/rustc_attr_parsing/src/attributes/rustc_internal.rs

@@ -1143,6 +1143,15 @@ impl<S: Stage> NoArgsAttributeParser<S> for RustcNonnullOptimizationGuaranteedPa
     const CREATE: fn(Span) -> AttributeKind = |_| AttributeKind::RustcNonnullOptimizationGuaranteed;
 }
 
+pub(crate) struct RustcPanicEntrypointParser;
+
+impl<S: Stage> NoArgsAttributeParser<S> for RustcPanicEntrypointParser {
+    const PATH: &[Symbol] = &[sym::rustc_panic_entrypoint];
+    const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::Error;
+    const ALLOWED_TARGETS: AllowedTargets = AllowedTargets::AllowList(&[Allow(Target::Fn)]);
+    const CREATE: fn(Span) -> AttributeKind = |_| AttributeKind::RustcPanicEntrypoint;
+}
+
 pub(crate) struct RustcStrictCoherenceParser;
 
 impl<S: Stage> NoArgsAttributeParser<S> for RustcStrictCoherenceParser {

compiler/rustc_attr_parsing/src/context.rs

@@ -318,6 +318,7 @@ attribute_parsers!(
         Single<WithoutArgs<RustcNonnullOptimizationGuaranteedParser>>,
         Single<WithoutArgs<RustcNounwindParser>>,
         Single<WithoutArgs<RustcOffloadKernelParser>>,
+        Single<WithoutArgs<RustcPanicEntrypointParser>>,
         Single<WithoutArgs<RustcParenSugarParser>>,
         Single<WithoutArgs<RustcPassByValueParser>>,
         Single<WithoutArgs<RustcPassIndirectlyInNonRusticAbisParser>>,

compiler/rustc_codegen_cranelift/src/base.rs

@@ -14,6 +14,7 @@ use rustc_middle::ty::layout::{FnAbiOf, HasTypingEnv as _};
 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);
@@ -1058,6 +1064,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

@@ -1,6 +1,6 @@
 use rustc_abi::{Align, ExternAbi};
 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};
@@ -15,7 +15,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::{
@@ -293,6 +293,9 @@ fn process_builtin_attrs(
                 codegen_fn_attrs.patchable_function_entry =
                     Some(PatchableFunctionEntry::from_prefix_and_entry(*prefix, *entry));
             }
+            AttributeKind::RustcPanicEntrypoint => {
+                codegen_fn_attrs.flags |= CodegenFnAttrFlags::PANIC_ENTRYPOINT
+            }
             _ => {}
         }
     }
@@ -385,6 +388,28 @@ fn apply_overrides(tcx: TyCtxt<'_>, did: LocalDefId, codegen_fn_attrs: &mut Code
             codegen_fn_attrs.flags |= CodegenFnAttrFlags::NO_MANGLE;
         }
     }
+
+    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::PANIC_ENTRYPOINT) {
+        // Panic entrypoints are always cold.
+        //
+        // If we have immediate-abort enabled, we want panic entrypoints to be inlined. We have a
+        // MIR transform that tries to patch out immediate-abort panic entrypoint calls, but it can
+        // miss indirect calls. In such cases, this should encourage the optimizer to clean up the
+        // mess that we missed.
+        //
+        // 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.
+        // FIXME: It seems unlikely that all choices of #[track_caller] were deliberate, but it can
+        // have a code size impact, so it makes sense to apply judiciously.
+        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;
+        }
+    }
 }
 
 #[derive(Diagnostic)]

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -14,6 +14,7 @@ use rustc_middle::{bug, span_bug};
 use rustc_session::config::OptLevel;
 use rustc_span::{Span, Spanned};
 use rustc_target::callconv::{ArgAbi, ArgAttributes, CastTarget, FnAbi, PassMode};
+use rustc_target::spec::PanicStrategy;
 use tracing::{debug, info};
 
 use super::operand::OperandRef;
@@ -175,22 +176,22 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
         if let Some(instance) = instance
             && is_call_from_compiler_builtins_to_upstream_monomorphization(tcx, instance)
         {
-            if destination.is_some() {
+            if instance.def.is_panic_entrypoint(tcx) {
+                info!(
+                    "compiler_builtins call to panic entrypoint {:?} replaced with abort",
+                    instance.def_id()
+                );
+                bx.abort();
+                bx.unreachable();
+                return MergingSucc::False;
+            } else {
                 let caller_def = fx.instance.def_id();
                 let e = CompilerBuiltinsCannotCall {
                     span: tcx.def_span(caller_def),
                     caller: with_no_trimmed_paths!(tcx.def_path_str(caller_def)),
                     callee: with_no_trimmed_paths!(tcx.def_path_str(instance.def_id())),
                 };
                 tcx.dcx().emit_err(e);
-            } else {
-                info!(
-                    "compiler_builtins call to diverging function {:?} replaced with abort",
-                    instance.def_id()
-                );
-                bx.abort();
-                bx.unreachable();
-                return MergingSucc::False;
             }
         }
 
@@ -736,6 +737,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

@@ -1323,6 +1323,11 @@ pub static BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         "`#[rustc_non_const_trait_method]` should only used by the standard library to mark trait methods \
         as non-const to allow large traits an easier transition to const"
     ),
+    rustc_attr!(
+        rustc_panic_entrypoint, AttributeType::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/attrs/data_structures.rs

@@ -1527,6 +1527,9 @@ pub enum AttributeKind {
     /// Represents `#[rustc_offload_kernel]`
     RustcOffloadKernel,
 
+    /// Represents `#[rustc_panic_entrypoint]`
+    RustcPanicEntrypoint,
+
     /// Represents `#[rustc_paren_sugar]`.
     RustcParenSugar(Span),
 

compiler/rustc_hir/src/attrs/encode_cross_crate.rs

@@ -170,6 +170,7 @@ impl AttributeKind {
             RustcObjcClass { .. } => No,
             RustcObjcSelector { .. } => No,
             RustcOffloadKernel => Yes,
+            RustcPanicEntrypoint => No,
             RustcParenSugar(..) => No,
             RustcPassByValue(..) => Yes,
             RustcPassIndirectlyInNonRusticAbis(..) => No,

compiler/rustc_hir/src/lang_items.rs

@@ -235,6 +235,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);