Rewrite implementation of #[alloc_error_handler]

The new implementation doesn't use weak lang items and instead changes
`#[alloc_error_handler]` to an attribute macro just like
`#[global_allocator]`.

The attribute will generate the `__rg_oom` function which is called by
the compiler-generated `__rust_alloc_error_handler`. If no `__rg_oom`
function is defined in any crate then the compiler shim will call
`__rdl_oom` in the alloc crate which will simply panic.

This also fixes link errors with `-C link-dead-code` with
`default_alloc_error_handler`: `__rg_oom` was previously defined in the
alloc crate and would attempt to reference the `oom` lang item, even if
it didn't exist. This worked as long as `__rg_oom` was excluded from
linking since it was not called.

This is a prerequisite for the stabilization of
`default_alloc_error_handler` (#102318).

compiler/rustc_builtin_macros/src/alloc_error_handler.rs

@@ -0,0 +1,104 @@
+use crate::util::check_builtin_macro_attribute;
+
+use rustc_ast::ptr::P;
+use rustc_ast::{self as ast, FnHeader, FnSig, Generics, StmtKind};
+use rustc_ast::{Fn, ItemKind, Stmt, TyKind, Unsafe};
+use rustc_expand::base::{Annotatable, ExtCtxt};
+use rustc_span::symbol::{kw, sym, Ident};
+use rustc_span::Span;
+use thin_vec::thin_vec;
+
+pub fn expand(
+    ecx: &mut ExtCtxt<'_>,
+    _span: Span,
+    meta_item: &ast::MetaItem,
+    item: Annotatable,
+) -> Vec<Annotatable> {
+    check_builtin_macro_attribute(ecx, meta_item, sym::alloc_error_handler);
+
+    let orig_item = item.clone();
+    let not_function = || {
+        ecx.sess
+            .parse_sess
+            .span_diagnostic
+            .span_err(item.span(), "alloc_error_handler must be a function");
+        vec![orig_item.clone()]
+    };
+
+    // Allow using `#[alloc_error_handler]` on an item statement
+    // FIXME - if we get deref patterns, use them to reduce duplication here
+    let (item, is_stmt, sig_span) = match &item {
+        Annotatable::Item(item) => match item.kind {
+            ItemKind::Fn(ref fn_kind) => (item, false, ecx.with_def_site_ctxt(fn_kind.sig.span)),
+            _ => return not_function(),
+        },
+        Annotatable::Stmt(stmt) => match &stmt.kind {
+            StmtKind::Item(item_) => match item_.kind {
+                ItemKind::Fn(ref fn_kind) => {
+                    (item_, true, ecx.with_def_site_ctxt(fn_kind.sig.span))
+                }
+                _ => return not_function(),
+            },
+            _ => return not_function(),
+        },
+        _ => return not_function(),
+    };
+
+    // Generate a bunch of new items using the AllocFnFactory
+    let span = ecx.with_def_site_ctxt(item.span);
+
+    // Generate item statements for the allocator methods.
+    let stmts = vec![generate_handler(ecx, item.ident, span, sig_span)];
+
+    // Generate anonymous constant serving as container for the allocator methods.
+    let const_ty = ecx.ty(sig_span, TyKind::Tup(Vec::new()));
+    let const_body = ecx.expr_block(ecx.block(span, stmts));
+    let const_item = ecx.item_const(span, Ident::new(kw::Underscore, span), const_ty, const_body);
+    let const_item = if is_stmt {
+        Annotatable::Stmt(P(ecx.stmt_item(span, const_item)))
+    } else {
+        Annotatable::Item(const_item)
+    };
+
+    // Return the original item and the new methods.
+    vec![orig_item, const_item]
+}
+
+// #[rustc_std_internal_symbol]
+// unsafe fn __rg_oom(size: usize, align: usize) -> ! {
+//     handler(core::alloc::Layout::from_size_align_unchecked(size, align))
+// }
+fn generate_handler(cx: &ExtCtxt<'_>, handler: Ident, span: Span, sig_span: Span) -> Stmt {
+    let usize = cx.path_ident(span, Ident::new(sym::usize, span));
+    let ty_usize = cx.ty_path(usize);
+    let size = Ident::from_str_and_span("size", span);
+    let align = Ident::from_str_and_span("align", span);
+
+    let layout_new = cx.std_path(&[sym::alloc, sym::Layout, sym::from_size_align_unchecked]);
+    let layout_new = cx.expr_path(cx.path(span, layout_new));
+    let layout =
+        cx.expr_call(span, layout_new, vec![cx.expr_ident(span, size), cx.expr_ident(span, align)]);
+
+    let call = cx.expr_call_ident(sig_span, handler, vec![layout]);
+
+    let never = ast::FnRetTy::Ty(cx.ty(span, TyKind::Never));
+    let params = vec![cx.param(span, size, ty_usize.clone()), cx.param(span, align, ty_usize)];
+    let decl = cx.fn_decl(params, never);
+    let header = FnHeader { unsafety: Unsafe::Yes(span), ..FnHeader::default() };
+    let sig = FnSig { decl, header, span: span };
+
+    let body = Some(cx.block_expr(call));
+    let kind = ItemKind::Fn(Box::new(Fn {
+        defaultness: ast::Defaultness::Final,
+        sig,
+        generics: Generics::default(),
+        body,
+    }));
+
+    let special = sym::rustc_std_internal_symbol;
+    let special = cx.meta_word(span, special);
+    let attrs = thin_vec![cx.attribute(special)];
+
+    let item = cx.item(span, Ident::from_str_and_span("__rg_oom", span), attrs, kind);
+    cx.stmt_item(sig_span, item)
+}

compiler/rustc_builtin_macros/src/lib.rs

@@ -25,6 +25,7 @@ use rustc_expand::base::{MacroExpanderFn, ResolverExpand, SyntaxExtensionKind};
 use rustc_expand::proc_macro::BangProcMacro;
 use rustc_span::symbol::sym;
 
+mod alloc_error_handler;
 mod assert;
 mod cfg;
 mod cfg_accessible;
@@ -94,6 +95,7 @@ pub fn register_builtin_macros(resolver: &mut dyn ResolverExpand) {
     }
 
     register_attr! {
+        alloc_error_handler: alloc_error_handler::expand,
         bench: test::expand_bench,
         cfg_accessible: cfg_accessible::Expander,
         cfg_eval: cfg_eval::expand,

compiler/rustc_codegen_cranelift/src/allocator.rs

@@ -5,6 +5,7 @@ use crate::prelude::*;
 
 use rustc_ast::expand::allocator::{AllocatorKind, AllocatorTy, ALLOCATOR_METHODS};
 use rustc_session::config::OomStrategy;
+use rustc_span::symbol::sym;
 
 /// Returns whether an allocator shim was created
 pub(crate) fn codegen(
@@ -23,7 +24,7 @@ pub(crate) fn codegen(
             module,
             unwind_context,
             kind,
-            tcx.lang_items().oom().is_some(),
+            tcx.alloc_error_handler_kind(()).unwrap(),
             tcx.sess.opts.unstable_opts.oom,
         );
         true
@@ -36,7 +37,7 @@ fn codegen_inner(
     module: &mut impl Module,
     unwind_context: &mut UnwindContext,
     kind: AllocatorKind,
-    has_alloc_error_handler: bool,
+    alloc_error_handler_kind: AllocatorKind,
     oom_strategy: OomStrategy,
 ) {
     let usize_ty = module.target_config().pointer_type();
@@ -108,12 +109,12 @@ fn codegen_inner(
         returns: vec![],
     };
 
-    let callee_name = if has_alloc_error_handler { "__rg_oom" } else { "__rdl_oom" };
+    let callee_name = alloc_error_handler_kind.fn_name(sym::oom);
 
     let func_id =
         module.declare_function("__rust_alloc_error_handler", Linkage::Export, &sig).unwrap();
 
-    let callee_func_id = module.declare_function(callee_name, Linkage::Import, &sig).unwrap();
+    let callee_func_id = module.declare_function(&callee_name, Linkage::Import, &sig).unwrap();
 
     let mut ctx = Context::new();
     ctx.func.signature = sig;

compiler/rustc_codegen_gcc/src/allocator.rs

@@ -7,7 +7,7 @@ use rustc_span::symbol::sym;
 
 use crate::GccContext;
 
-pub(crate) unsafe fn codegen(tcx: TyCtxt<'_>, mods: &mut GccContext, _module_name: &str, kind: AllocatorKind, has_alloc_error_handler: bool) {
+pub(crate) unsafe fn codegen(tcx: TyCtxt<'_>, mods: &mut GccContext, _module_name: &str, kind: AllocatorKind, alloc_error_handler_kind: AllocatorKind) {
     let context = &mods.context;
     let usize =
         match tcx.sess.target.pointer_width {
@@ -90,14 +90,7 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt<'_>, mods: &mut GccContext, _module_nam
         .collect();
     let func = context.new_function(None, FunctionType::Exported, void, &args, name, false);
 
-    let kind =
+    let callee = alloc_error_handler_kind.fn_name(sym::oom);
-        if has_alloc_error_handler {
-            AllocatorKind::Global
-        }
-        else {
-            AllocatorKind::Default
-        };
-    let callee = kind.fn_name(sym::oom);
     let args: Vec<_> = types.iter().enumerate()
         .map(|(index, typ)| context.new_parameter(None, *typ, &format!("param{}", index)))
         .collect();

compiler/rustc_codegen_gcc/src/lib.rs

@@ -153,11 +153,11 @@ impl CodegenBackend for GccCodegenBackend {
 }
 
 impl ExtraBackendMethods for GccCodegenBackend {
-    fn codegen_allocator<'tcx>(&self, tcx: TyCtxt<'tcx>, module_name: &str, kind: AllocatorKind, has_alloc_error_handler: bool) -> Self::Module {
+    fn codegen_allocator<'tcx>(&self, tcx: TyCtxt<'tcx>, module_name: &str, kind: AllocatorKind, alloc_error_handler_kind: AllocatorKind) -> Self::Module {
         let mut mods = GccContext {
             context: Context::default(),
         };
-        unsafe { allocator::codegen(tcx, &mut mods, module_name, kind, has_alloc_error_handler); }
+        unsafe { allocator::codegen(tcx, &mut mods, module_name, kind, alloc_error_handler_kind); }
         mods
     }
 

compiler/rustc_codegen_llvm/src/allocator.rs

@@ -15,7 +15,7 @@ pub(crate) unsafe fn codegen(
     module_llvm: &mut ModuleLlvm,
     module_name: &str,
     kind: AllocatorKind,
-    has_alloc_error_handler: bool,
+    alloc_error_handler_kind: AllocatorKind,
 ) {
     let llcx = &*module_llvm.llcx;
     let llmod = module_llvm.llmod();
@@ -117,8 +117,7 @@ pub(crate) unsafe fn codegen(
         attributes::apply_to_llfn(llfn, llvm::AttributePlace::Function, &[uwtable]);
     }
 
-    let kind = if has_alloc_error_handler { AllocatorKind::Global } else { AllocatorKind::Default };
+    let callee = alloc_error_handler_kind.fn_name(sym::oom);
-    let callee = kind.fn_name(sym::oom);
     let callee = llvm::LLVMRustGetOrInsertFunction(llmod, callee.as_ptr().cast(), callee.len(), ty);
     // -> ! DIFlagNoReturn
     attributes::apply_to_llfn(callee, llvm::AttributePlace::Function, &[no_return]);

compiler/rustc_codegen_llvm/src/lib.rs

@@ -108,11 +108,11 @@ impl ExtraBackendMethods for LlvmCodegenBackend {
         tcx: TyCtxt<'tcx>,
         module_name: &str,
         kind: AllocatorKind,
-        has_alloc_error_handler: bool,
+        alloc_error_handler_kind: AllocatorKind,
     ) -> ModuleLlvm {
         let mut module_llvm = ModuleLlvm::new_metadata(tcx, module_name);
         unsafe {
-            allocator::codegen(tcx, &mut module_llvm, module_name, kind, has_alloc_error_handler);
+            allocator::codegen(tcx, &mut module_llvm, module_name, kind, alloc_error_handler_kind);
         }
         module_llvm
     }

compiler/rustc_codegen_ssa/src/back/symbol_export.rs

@@ -193,8 +193,11 @@ fn exported_symbols_provider_local<'tcx>(
     }
 
     if tcx.allocator_kind(()).is_some() {
-        for method in ALLOCATOR_METHODS {
+        for symbol_name in ALLOCATOR_METHODS
-            let symbol_name = format!("__rust_{}", method.name);
+            .iter()
+            .map(|method| format!("__rust_{}", method.name))
+            .chain(["__rust_alloc_error_handler".to_string(), OomStrategy::SYMBOL.to_string()])
+        {
             let exported_symbol = ExportedSymbol::NoDefId(SymbolName::new(tcx, &symbol_name));
 
             symbols.push((

compiler/rustc_codegen_ssa/src/base.rs

@@ -638,7 +638,14 @@ pub fn codegen_crate<B: ExtraBackendMethods>(
         let llmod_id =
             cgu_name_builder.build_cgu_name(LOCAL_CRATE, &["crate"], Some("allocator")).to_string();
         let module_llvm = tcx.sess.time("write_allocator_module", || {
-            backend.codegen_allocator(tcx, &llmod_id, kind, tcx.lang_items().oom().is_some())
+            backend.codegen_allocator(
+                tcx,
+                &llmod_id,
+                kind,
+                // If allocator_kind is Some then alloc_error_handler_kind must
+                // also be Some.
+                tcx.alloc_error_handler_kind(()).unwrap(),
+            )
         });
 
         Some(ModuleCodegen { name: llmod_id, module_llvm, kind: ModuleKind::Allocator })

compiler/rustc_codegen_ssa/src/traits/backend.rs

@@ -119,7 +119,7 @@ pub trait ExtraBackendMethods: CodegenBackend + WriteBackendMethods + Sized + Se
         tcx: TyCtxt<'tcx>,
         module_name: &str,
         kind: AllocatorKind,
-        has_alloc_error_handler: bool,
+        alloc_error_handler_kind: AllocatorKind,
     ) -> Self::Module;
     /// This generates the codegen unit and returns it along with
     /// a `u64` giving an estimate of the unit's processing cost.

compiler/rustc_error_messages/locales/en-US/metadata.ftl

@@ -150,12 +150,28 @@ metadata_no_multiple_global_alloc =
 metadata_prev_global_alloc =
     previous global allocator defined here
 
+metadata_no_multiple_alloc_error_handler =
+    cannot define multiple allocation error handlers
+    .label = cannot define a new allocation error handler
+
+metadata_prev_alloc_error_handler =
+    previous allocation error handler defined here
+
 metadata_conflicting_global_alloc =
     the `#[global_allocator]` in {$other_crate_name} conflicts with global allocator in: {$crate_name}
 
+metadata_conflicting_alloc_error_handler =
+    the `#[alloc_error_handler]` in {$other_crate_name} conflicts with allocation error handler in: {$crate_name}
+
 metadata_global_alloc_required =
     no global memory allocator found but one is required; link to std or add `#[global_allocator]` to a static item that implements the GlobalAlloc trait
 
+metadata_alloc_func_required =
+    `#[alloc_error_handler]` function required, but not found
+
+metadata_missing_alloc_error_handler =
+    use `#![feature(default_alloc_error_handler)]` for a default error handler
+
 metadata_no_transitive_needs_dep =
     the crate `{$crate_name}` cannot depend on a crate that needs {$needs_crate_name}, but it depends on `{$deps_crate_name}`
 

compiler/rustc_error_messages/locales/en-US/passes.ftl

@@ -367,12 +367,6 @@ passes_unknown_external_lang_item =
 passes_missing_panic_handler =
     `#[panic_handler]` function required, but not found
 
-passes_alloc_func_required =
-    `#[alloc_error_handler]` function required, but not found
-
-passes_missing_alloc_error_handler =
-    use `#![feature(default_alloc_error_handler)]` for a default error handler
-
 passes_missing_lang_item =
     language item required, but not found: `{$name}`
     .note = this can occur when a binary crate with `#![no_std]` is compiled for a target where `{$name}` is defined in the standard library

compiler/rustc_feature/src/builtin_attrs.rs

@@ -554,10 +554,6 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
     rustc_attr!(rustc_reallocator, Normal, template!(Word), WarnFollowing, IMPL_DETAIL),
     rustc_attr!(rustc_deallocator, Normal, template!(Word), WarnFollowing, IMPL_DETAIL),
     rustc_attr!(rustc_allocator_zeroed, Normal, template!(Word), WarnFollowing, IMPL_DETAIL),
-    gated!(
-        alloc_error_handler, Normal, template!(Word), WarnFollowing,
-        experimental!(alloc_error_handler)
-    ),
     gated!(
         default_lib_allocator, Normal, template!(Word), WarnFollowing, allocator_internals,
         experimental!(default_lib_allocator),

compiler/rustc_hir/src/lang_items.rs

@@ -126,14 +126,12 @@ impl<CTX> HashStable<CTX> for LangItem {
 }
 
 /// Extracts the first `lang = "$name"` out of a list of attributes.
-/// The attributes `#[panic_handler]` and `#[alloc_error_handler]`
+/// The `#[panic_handler]` attribute is also extracted out when found.
-/// are also extracted out when found.
 pub fn extract(attrs: &[ast::Attribute]) -> Option<(Symbol, Span)> {
     attrs.iter().find_map(|attr| {
         Some(match attr {
             _ if attr.has_name(sym::lang) => (attr.value_str()?, attr.span),
             _ if attr.has_name(sym::panic_handler) => (sym::panic_impl, attr.span),
-            _ if attr.has_name(sym::alloc_error_handler) => (sym::oom, attr.span),
             _ => return None,
         })
     })
@@ -240,7 +238,6 @@ language_item_table! {
     ExchangeMalloc,          sym::exchange_malloc,     exchange_malloc_fn,         Target::Fn,             GenericRequirement::None;
     BoxFree,                 sym::box_free,            box_free_fn,                Target::Fn,             GenericRequirement::Minimum(1);
     DropInPlace,             sym::drop_in_place,       drop_in_place_fn,           Target::Fn,             GenericRequirement::Minimum(1);
-    Oom,                     sym::oom,                 oom,                        Target::Fn,             GenericRequirement::None;
     AllocLayout,             sym::alloc_layout,        alloc_layout,               Target::Struct,         GenericRequirement::None;
 
     Start,                   sym::start,               start_fn,                   Target::Fn,             GenericRequirement::Exact(1);

compiler/rustc_hir/src/weak_lang_items.rs

@@ -27,5 +27,4 @@ weak_lang_items! {
     PanicImpl,          rust_begin_unwind;
     EhPersonality,      rust_eh_personality;
     EhCatchTypeinfo,    rust_eh_catch_typeinfo;
-    Oom,                rust_oom;
 }

compiler/rustc_hir_typeck/src/check.rs

@@ -211,13 +211,6 @@ pub(super) fn check_fn<'a, 'tcx>(
         check_panic_info_fn(tcx, panic_impl_did.expect_local(), fn_sig, decl, declared_ret_ty);
     }
 
-    // Check that a function marked as `#[alloc_error_handler]` has signature `fn(Layout) -> !`
-    if let Some(alloc_error_handler_did) = tcx.lang_items().oom()
-        && alloc_error_handler_did == hir.local_def_id(fn_id).to_def_id()
-    {
-        check_alloc_error_fn(tcx, alloc_error_handler_did.expect_local(), fn_sig, decl, declared_ret_ty);
-    }
-
     (fcx, gen_ty)
 }
 
@@ -273,52 +266,3 @@ fn check_panic_info_fn(
         tcx.sess.span_err(span, "should have no const parameters");
     }
 }
-
-fn check_alloc_error_fn(
-    tcx: TyCtxt<'_>,
-    fn_id: LocalDefId,
-    fn_sig: ty::FnSig<'_>,
-    decl: &hir::FnDecl<'_>,
-    declared_ret_ty: Ty<'_>,
-) {
-    let Some(alloc_layout_did) = tcx.lang_items().alloc_layout() else {
-        tcx.sess.err("language item required, but not found: `alloc_layout`");
-        return;
-    };
-
-    if *declared_ret_ty.kind() != ty::Never {
-        tcx.sess.span_err(decl.output.span(), "return type should be `!`");
-    }
-
-    let inputs = fn_sig.inputs();
-    if inputs.len() != 1 {
-        tcx.sess.span_err(tcx.def_span(fn_id), "function should have one argument");
-        return;
-    }
-
-    let arg_is_alloc_layout = match inputs[0].kind() {
-        ty::Adt(ref adt, _) => adt.did() == alloc_layout_did,
-        _ => false,
-    };
-
-    if !arg_is_alloc_layout {
-        tcx.sess.span_err(decl.inputs[0].span, "argument should be `Layout`");
-    }
-
-    let DefKind::Fn = tcx.def_kind(fn_id) else {
-        let span = tcx.def_span(fn_id);
-        tcx.sess.span_err(span, "`#[alloc_error_handler]` should be a function");
-        return;
-    };
-
-    let generic_counts = tcx.generics_of(fn_id).own_counts();
-    if generic_counts.types != 0 {
-        let span = tcx.def_span(fn_id);
-        tcx.sess.span_err(span, "`#[alloc_error_handler]` function should have no type parameters");
-    }
-    if generic_counts.consts != 0 {
-        let span = tcx.def_span(fn_id);
-        tcx.sess
-            .span_err(span, "`#[alloc_error_handler]` function should have no const parameters");
-    }
-}

compiler/rustc_metadata/src/creader.rs

@@ -1,8 +1,10 @@
 //! Validates all used crates and extern libraries and loads their metadata
 
 use crate::errors::{
-    ConflictingGlobalAlloc, CrateNotPanicRuntime, GlobalAllocRequired, NoMultipleGlobalAlloc,
+    AllocFuncRequired, ConflictingAllocErrorHandler, ConflictingGlobalAlloc, CrateNotPanicRuntime,
-    NoPanicStrategy, NoTransitiveNeedsDep, NotProfilerRuntime, ProfilerBuiltinsNeedsCore,
+    GlobalAllocRequired, MissingAllocErrorHandler, NoMultipleAllocErrorHandler,
+    NoMultipleGlobalAlloc, NoPanicStrategy, NoTransitiveNeedsDep, NotProfilerRuntime,
+    ProfilerBuiltinsNeedsCore,
 };
 use crate::locator::{CrateError, CrateLocator, CratePaths};
 use crate::rmeta::{CrateDep, CrateMetadata, CrateNumMap, CrateRoot, MetadataBlob};
@@ -41,8 +43,13 @@ pub struct CStore {
     /// This crate needs an allocator and either provides it itself, or finds it in a dependency.
     /// If the above is true, then this field denotes the kind of the found allocator.
     allocator_kind: Option<AllocatorKind>,
+    /// This crate needs an allocation error handler and either provides it itself, or finds it in a dependency.
+    /// If the above is true, then this field denotes the kind of the found allocator.
+    alloc_error_handler_kind: Option<AllocatorKind>,
     /// This crate has a `#[global_allocator]` item.
     has_global_allocator: bool,
+    /// This crate has a `#[alloc_error_handler]` item.
+    has_alloc_error_handler: bool,
 
     /// This map is used to verify we get no hash conflicts between
     /// `StableCrateId` values.
@@ -197,10 +204,18 @@ impl CStore {
         self.allocator_kind
     }
 
+    pub(crate) fn alloc_error_handler_kind(&self) -> Option<AllocatorKind> {
+        self.alloc_error_handler_kind
+    }
+
     pub(crate) fn has_global_allocator(&self) -> bool {
         self.has_global_allocator
     }
 
+    pub(crate) fn has_alloc_error_handler(&self) -> bool {
+        self.has_alloc_error_handler
+    }
+
     pub fn report_unused_deps(&self, tcx: TyCtxt<'_>) {
         let json_unused_externs = tcx.sess.opts.json_unused_externs;
 
@@ -247,7 +262,9 @@ impl<'a> CrateLoader<'a> {
                 metas: IndexVec::from_elem_n(None, 1),
                 injected_panic_runtime: None,
                 allocator_kind: None,
+                alloc_error_handler_kind: None,
                 has_global_allocator: false,
+                has_alloc_error_handler: false,
                 stable_crate_ids,
                 unused_externs: Vec::new(),
             },
@@ -792,6 +809,13 @@ impl<'a> CrateLoader<'a> {
             }
             spans => !spans.is_empty(),
         };
+        self.cstore.has_alloc_error_handler = match &*alloc_error_handler_spans(&self.sess, krate) {
+            [span1, span2, ..] => {
+                self.sess.emit_err(NoMultipleAllocErrorHandler { span2: *span2, span1: *span1 });
+                true
+            }
+            spans => !spans.is_empty(),
+        };
 
         // Check to see if we actually need an allocator. This desire comes
         // about through the `#![needs_allocator]` attribute and is typically
@@ -832,22 +856,48 @@ impl<'a> CrateLoader<'a> {
                 }
             }
         }
+        let mut alloc_error_handler =
+            self.cstore.has_alloc_error_handler.then(|| Symbol::intern("this crate"));
+        for (_, data) in self.cstore.iter_crate_data() {
+            if data.has_alloc_error_handler() {
+                match alloc_error_handler {
+                    Some(other_crate) => {
+                        self.sess.emit_err(ConflictingAllocErrorHandler {
+                            crate_name: data.name(),
+                            other_crate_name: other_crate,
+                        });
+                    }
+                    None => alloc_error_handler = Some(data.name()),
+                }
+            }
+        }
 
         if global_allocator.is_some() {
             self.cstore.allocator_kind = Some(AllocatorKind::Global);
-            return;
+        } else {
+            // Ok we haven't found a global allocator but we still need an
+            // allocator. At this point our allocator request is typically fulfilled
+            // by the standard library, denoted by the `#![default_lib_allocator]`
+            // attribute.
+            if !self.sess.contains_name(&krate.attrs, sym::default_lib_allocator)
+                && !self.cstore.iter_crate_data().any(|(_, data)| data.has_default_lib_allocator())
+            {
+                self.sess.emit_err(GlobalAllocRequired);
+            }
+            self.cstore.allocator_kind = Some(AllocatorKind::Default);
         }
 
-        // Ok we haven't found a global allocator but we still need an
+        if alloc_error_handler.is_some() {
-        // allocator. At this point our allocator request is typically fulfilled
+            self.cstore.alloc_error_handler_kind = Some(AllocatorKind::Global);
-        // by the standard library, denoted by the `#![default_lib_allocator]`
+        } else {
-        // attribute.
+            // The alloc crate provides a default allocation error handler if
-        if !self.sess.contains_name(&krate.attrs, sym::default_lib_allocator)
+            // one isn't specified.
-            && !self.cstore.iter_crate_data().any(|(_, data)| data.has_default_lib_allocator())
+            if !self.sess.features_untracked().default_alloc_error_handler {
-        {
+                self.sess.emit_err(AllocFuncRequired);
-            self.sess.emit_err(GlobalAllocRequired);
+                self.sess.emit_note(MissingAllocErrorHandler);
+            }
+            self.cstore.alloc_error_handler_kind = Some(AllocatorKind::Default);
         }
-        self.cstore.allocator_kind = Some(AllocatorKind::Default);
     }
 
     fn inject_dependency_if(
@@ -1023,3 +1073,26 @@ fn global_allocator_spans(sess: &Session, krate: &ast::Crate) -> Vec<Span> {
     visit::walk_crate(&mut f, krate);
     f.spans
 }
+
+fn alloc_error_handler_spans(sess: &Session, krate: &ast::Crate) -> Vec<Span> {
+    struct Finder<'a> {
+        sess: &'a Session,
+        name: Symbol,
+        spans: Vec<Span>,
+    }
+    impl<'ast, 'a> visit::Visitor<'ast> for Finder<'a> {
+        fn visit_item(&mut self, item: &'ast ast::Item) {
+            if item.ident.name == self.name
+                && self.sess.contains_name(&item.attrs, sym::rustc_std_internal_symbol)
+            {
+                self.spans.push(item.span);
+            }
+            visit::walk_item(self, item)
+        }
+    }
+
+    let name = Symbol::intern(&AllocatorKind::Global.fn_name(sym::oom));
+    let mut f = Finder { sess, name, spans: Vec::new() };
+    visit::walk_crate(&mut f, krate);
+    f.spans
+}

compiler/rustc_metadata/src/errors.rs

@@ -343,6 +343,16 @@ pub struct NoMultipleGlobalAlloc {
     pub span1: Span,
 }
 
+#[derive(Diagnostic)]
+#[diag(metadata_no_multiple_alloc_error_handler)]
+pub struct NoMultipleAllocErrorHandler {
+    #[primary_span]
+    #[label]
+    pub span2: Span,
+    #[label(metadata_prev_alloc_error_handler)]
+    pub span1: Span,
+}
+
 #[derive(Diagnostic)]
 #[diag(metadata_conflicting_global_alloc)]
 pub struct ConflictingGlobalAlloc {
@@ -350,10 +360,25 @@ pub struct ConflictingGlobalAlloc {
     pub other_crate_name: Symbol,
 }
 
+#[derive(Diagnostic)]
+#[diag(metadata_conflicting_alloc_error_handler)]
+pub struct ConflictingAllocErrorHandler {
+    pub crate_name: Symbol,
+    pub other_crate_name: Symbol,
+}
+
 #[derive(Diagnostic)]
 #[diag(metadata_global_alloc_required)]
 pub struct GlobalAllocRequired;
 
+#[derive(Diagnostic)]
+#[diag(metadata_alloc_func_required)]
+pub struct AllocFuncRequired;
+
+#[derive(Diagnostic)]
+#[diag(metadata_missing_alloc_error_handler)]
+pub struct MissingAllocErrorHandler;
+
 #[derive(Diagnostic)]
 #[diag(metadata_no_transitive_needs_dep)]
 pub struct NoTransitiveNeedsDep<'a> {

compiler/rustc_metadata/src/rmeta/decoder.rs

@@ -1764,6 +1764,10 @@ impl CrateMetadata {
         self.root.has_global_allocator
     }
 
+    pub(crate) fn has_alloc_error_handler(&self) -> bool {
+        self.root.has_alloc_error_handler
+    }
+
     pub(crate) fn has_default_lib_allocator(&self) -> bool {
         self.root.has_default_lib_allocator
     }

compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs

@@ -255,6 +255,7 @@ provide! { tcx, def_id, other, cdata,
     is_panic_runtime => { cdata.root.panic_runtime }
     is_compiler_builtins => { cdata.root.compiler_builtins }
     has_global_allocator => { cdata.root.has_global_allocator }
+    has_alloc_error_handler => { cdata.root.has_alloc_error_handler }
     has_panic_handler => { cdata.root.has_panic_handler }
     is_profiler_runtime => { cdata.root.profiler_runtime }
     required_panic_strategy => { cdata.root.required_panic_strategy }
@@ -339,6 +340,7 @@ pub(in crate::rmeta) fn provide(providers: &mut Providers) {
     // resolve! Does this work? Unsure! That's what the issue is about
     *providers = Providers {
         allocator_kind: |tcx, ()| CStore::from_tcx(tcx).allocator_kind(),
+        alloc_error_handler_kind: |tcx, ()| CStore::from_tcx(tcx).alloc_error_handler_kind(),
         is_private_dep: |_tcx, cnum| {
             assert_eq!(cnum, LOCAL_CRATE);
             false
@@ -464,6 +466,10 @@ pub(in crate::rmeta) fn provide(providers: &mut Providers) {
             assert_eq!(cnum, LOCAL_CRATE);
             CStore::from_tcx(tcx).has_global_allocator()
         },
+        has_alloc_error_handler: |tcx, cnum| {
+            assert_eq!(cnum, LOCAL_CRATE);
+            CStore::from_tcx(tcx).has_alloc_error_handler()
+        },
         postorder_cnums: |tcx, ()| {
             tcx.arena
                 .alloc_slice(&CStore::from_tcx(tcx).crate_dependencies_in_postorder(LOCAL_CRATE))

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -670,6 +670,7 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
                 panic_in_drop_strategy: tcx.sess.opts.unstable_opts.panic_in_drop,
                 edition: tcx.sess.edition(),
                 has_global_allocator: tcx.has_global_allocator(LOCAL_CRATE),
+                has_alloc_error_handler: tcx.has_alloc_error_handler(LOCAL_CRATE),
                 has_panic_handler: tcx.has_panic_handler(LOCAL_CRATE),
                 has_default_lib_allocator: tcx
                     .sess

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -223,6 +223,7 @@ pub(crate) struct CrateRoot {
     panic_in_drop_strategy: PanicStrategy,
     edition: Edition,
     has_global_allocator: bool,
+    has_alloc_error_handler: bool,
     has_panic_handler: bool,
     has_default_lib_allocator: bool,
 

compiler/rustc_middle/src/query/mod.rs

@@ -1391,6 +1391,13 @@ rustc_queries! {
         desc { "checking if the crate has_global_allocator" }
         separate_provide_extern
     }
+    query has_alloc_error_handler(_: CrateNum) -> bool {
+        // This query depends on untracked global state in CStore
+        eval_always
+        fatal_cycle
+        desc { "checking if the crate has_alloc_error_handler" }
+        separate_provide_extern
+    }
     query has_panic_handler(_: CrateNum) -> bool {
         fatal_cycle
         desc { "checking if the crate has_panic_handler" }
@@ -1761,6 +1768,10 @@ rustc_queries! {
         eval_always
         desc { "getting the allocator kind for the current crate" }
     }
+    query alloc_error_handler_kind(_: ()) -> Option<AllocatorKind> {
+        eval_always
+        desc { "alloc error handler kind for the current crate" }
+    }
 
     query upvars_mentioned(def_id: DefId) -> Option<&'tcx FxIndexMap<hir::HirId, hir::Upvar>> {
         desc { |tcx| "collecting upvars mentioned in `{}`", tcx.def_path_str(def_id) }

compiler/rustc_passes/src/dead.rs

@@ -470,11 +470,6 @@ fn has_allow_dead_code_or_lang_attr_helper(
         return true;
     }
 
-    // (To be) stable attribute for #[lang = "oom"]
-    if tcx.sess.contains_name(attrs, sym::alloc_error_handler) {
-        return true;
-    }
-
     let def_id = tcx.hir().local_def_id(id);
     if tcx.def_kind(def_id).has_codegen_attrs() {
         let cg_attrs = tcx.codegen_fn_attrs(def_id);

compiler/rustc_passes/src/errors.rs

@@ -701,14 +701,6 @@ pub struct UnknownExternLangItem {
 #[diag(passes_missing_panic_handler)]
 pub struct MissingPanicHandler;
 
-#[derive(Diagnostic)]
-#[diag(passes_alloc_func_required)]
-pub struct AllocFuncRequired;
-
-#[derive(Diagnostic)]
-#[diag(passes_missing_alloc_error_handler)]
-pub struct MissingAllocErrorHandler;
-
 #[derive(Diagnostic)]
 #[diag(passes_missing_lang_item)]
 #[note]

compiler/rustc_passes/src/weak_lang_items.rs

@@ -7,10 +7,7 @@ use rustc_middle::middle::lang_items::required;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::config::CrateType;
 
-use crate::errors::{
+use crate::errors::{MissingLangItem, MissingPanicHandler, UnknownExternLangItem};
-    AllocFuncRequired, MissingAllocErrorHandler, MissingLangItem, MissingPanicHandler,
-    UnknownExternLangItem,
-};
 
 /// Checks the crate for usage of weak lang items, returning a vector of all the
 /// language items required by this crate, but not defined yet.
@@ -69,11 +66,6 @@ fn verify<'tcx>(tcx: TyCtxt<'tcx>, items: &lang_items::LanguageItems) {
         if missing.contains(&item) && required(tcx, item) && items.get(item).is_none() {
             if item == LangItem::PanicImpl {
                 tcx.sess.emit_err(MissingPanicHandler);
-            } else if item == LangItem::Oom {
-                if !tcx.features().default_alloc_error_handler {
-                    tcx.sess.emit_err(AllocFuncRequired);
-                    tcx.sess.emit_note(MissingAllocErrorHandler);
-                }
             } else {
                 tcx.sess.emit_err(MissingLangItem { name: item.name() });
             }

library/alloc/src/alloc.rs

@@ -402,19 +402,18 @@ pub use std::alloc::handle_alloc_error;
 #[allow(unused_attributes)]
 #[unstable(feature = "alloc_internals", issue = "none")]
 pub mod __alloc_error_handler {
-    use crate::alloc::Layout;
+    // called via generated `__rust_alloc_error_handler` if there is no
-
+    // `#[alloc_error_handler]`.
-    // called via generated `__rust_alloc_error_handler`
-
-    // if there is no `#[alloc_error_handler]`
     #[rustc_std_internal_symbol]
     pub unsafe fn __rdl_oom(size: usize, _align: usize) -> ! {
         panic!("memory allocation of {size} bytes failed")
     }
 
-    // if there is an `#[alloc_error_handler]`
+    #[cfg(bootstrap)]
     #[rustc_std_internal_symbol]
     pub unsafe fn __rg_oom(size: usize, align: usize) -> ! {
+        use crate::alloc::Layout;
+
         let layout = unsafe { Layout::from_size_align_unchecked(size, align) };
         extern "Rust" {
             #[lang = "oom"]

library/core/src/macros/mod.rs

@@ -1511,6 +1511,17 @@ pub(crate) mod builtin {
         /* compiler built-in */
     }
 
+    /// Attribute macro applied to a function to register it as a handler for allocation failure.
+    ///
+    /// See also [`std::alloc::handle_alloc_error`](../../../std/alloc/fn.handle_alloc_error.html).
+    #[cfg(not(bootstrap))]
+    #[unstable(feature = "alloc_error_handler", issue = "51540")]
+    #[allow_internal_unstable(rustc_attrs)]
+    #[rustc_builtin_macro]
+    pub macro alloc_error_handler($item:item) {
+        /* compiler built-in */
+    }
+
     /// Keeps the item it's applied to if the passed path is accessible, and removes it otherwise.
     #[unstable(
         feature = "cfg_accessible",

library/core/src/prelude/v1.rs

@@ -75,6 +75,9 @@ pub use crate::macros::builtin::{RustcDecodable, RustcEncodable};
 
 // Do not `doc(no_inline)` so that they become doc items on their own
 // (no public module for them to be re-exported from).
+#[cfg(not(bootstrap))]
+#[stable(feature = "builtin_macro_prelude", since = "1.38.0")]
+pub use crate::macros::builtin::alloc_error_handler;
 #[stable(feature = "builtin_macro_prelude", since = "1.38.0")]
 pub use crate::macros::builtin::{bench, derive, global_allocator, test, test_case};
 

library/std/src/prelude/v1.rs

@@ -59,6 +59,9 @@ pub use core::prelude::v1::{RustcDecodable, RustcEncodable};
 
 // Do not `doc(no_inline)` so that they become doc items on their own
 // (no public module for them to be re-exported from).
+#[cfg(not(bootstrap))]
+#[stable(feature = "builtin_macro_prelude", since = "1.38.0")]
+pub use core::prelude::v1::alloc_error_handler;
 #[stable(feature = "builtin_macro_prelude", since = "1.38.0")]
 pub use core::prelude::v1::{bench, derive, global_allocator, test, test_case};
 

src/test/run-make-fulldeps/issue-51671/Makefile

@@ -6,4 +6,4 @@ all:
 	$(RUSTC) --emit=obj app.rs
 	nm $(TMPDIR)/app.o | $(CGREP) rust_begin_unwind
 	nm $(TMPDIR)/app.o | $(CGREP) rust_eh_personality
-	nm $(TMPDIR)/app.o | $(CGREP) rust_oom
+	nm $(TMPDIR)/app.o | $(CGREP) __rg_oom

src/test/run-make-fulldeps/issue-51671/app.rs

@@ -1,5 +1,5 @@
 #![crate_type = "bin"]
-#![feature(lang_items)]
+#![feature(lang_items, alloc_error_handler)]
 #![no_main]
 #![no_std]
 
@@ -14,7 +14,7 @@ fn panic(_: &PanicInfo) -> ! {
 #[lang = "eh_personality"]
 fn eh() {}
 
-#[lang = "oom"]
+#[alloc_error_handler]
 fn oom(_: Layout) -> ! {
     loop {}
 }

src/test/ui/alloc-error/alloc-error-handler-bad-signature-1.rs

@@ -8,8 +8,8 @@ use core::alloc::Layout;
 
 #[alloc_error_handler]
 fn oom(
-    info: &Layout, //~ ERROR argument should be `Layout`
+    info: &Layout, //~^ ERROR mismatched types
-) -> () //~ ERROR return type should be `!`
+) -> () //~^^ ERROR mismatched types
 {
     loop {}
 }

src/test/ui/alloc-error/alloc-error-handler-bad-signature-1.stderr

@@ -1,14 +1,50 @@
-error: return type should be `!`
+error[E0308]: mismatched types
-  --> $DIR/alloc-error-handler-bad-signature-1.rs:12:6
+  --> $DIR/alloc-error-handler-bad-signature-1.rs:10:1
    |
-LL | ) -> ()
+LL |    #[alloc_error_handler]
-   |      ^^
+   |    ---------------------- in this procedural macro expansion
-
+LL |    fn oom(
-error: argument should be `Layout`
+   |   _^
-  --> $DIR/alloc-error-handler-bad-signature-1.rs:11:11
+   |  |_|
+   | ||
+LL | ||     info: &Layout,
+LL | || ) -> ()
+   | ||_______- arguments to this function are incorrect
+LL | |  {
+LL | |      loop {}
+LL | |  }
+   | |__^ expected `&Layout`, found struct `Layout`
    |
+note: function defined here
+  --> $DIR/alloc-error-handler-bad-signature-1.rs:10:4
+   |
+LL | fn oom(
+   |    ^^^
 LL |     info: &Layout,
-   |           ^^^^^^^
+   |     -------------
+   = note: this error originates in the attribute macro `alloc_error_handler` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error[E0308]: mismatched types
+  --> $DIR/alloc-error-handler-bad-signature-1.rs:10:1
+   |
+LL |    #[alloc_error_handler]
+   |    ---------------------- in this procedural macro expansion
+LL |    fn oom(
+   |   _^
+   |  |_|
+   | ||
+LL | ||     info: &Layout,
+LL | || ) -> ()
+   | ||_______^ expected `!`, found `()`
+LL | |  {
+LL | |      loop {}
+LL | |  }
+   | |__- expected `!` because of return type
+   |
+   = note:   expected type `!`
+           found unit type `()`
+   = note: this error originates in the attribute macro `alloc_error_handler` (in Nightly builds, run with -Z macro-backtrace for more info)
 
 error: aborting due to 2 previous errors
 
+For more information about this error, try `rustc --explain E0308`.

src/test/ui/alloc-error/alloc-error-handler-bad-signature-2.rs

@@ -8,8 +8,8 @@ struct Layout;
 
 #[alloc_error_handler]
 fn oom(
-    info: Layout, //~ ERROR argument should be `Layout`
+    info: Layout, //~^ ERROR mismatched types
-) { //~ ERROR return type should be `!`
+) { //~^^ ERROR mismatched types
     loop {}
 }
 

src/test/ui/alloc-error/alloc-error-handler-bad-signature-2.stderr

@@ -1,14 +1,63 @@
-error: return type should be `!`
+error[E0308]: mismatched types
-  --> $DIR/alloc-error-handler-bad-signature-2.rs:12:3
+  --> $DIR/alloc-error-handler-bad-signature-2.rs:10:1
    |
-LL | ) {
+LL |    #[alloc_error_handler]
-   |   ^
+   |    ---------------------- in this procedural macro expansion
-
+LL |    fn oom(
-error: argument should be `Layout`
+   |   _^
-  --> $DIR/alloc-error-handler-bad-signature-2.rs:11:11
+   |  |_|
+   | ||
+LL | ||     info: Layout,
+LL | || ) {
+   | || -
+   | ||_|
+   | |  arguments to this function are incorrect
+LL | |      loop {}
+LL | |  }
+   | |__^ expected struct `Layout`, found struct `core::alloc::Layout`
    |
+   = note: struct `core::alloc::Layout` and struct `Layout` have similar names, but are actually distinct types
+note: struct `core::alloc::Layout` is defined in crate `core`
+  --> $SRC_DIR/core/src/alloc/layout.rs:LL:COL
+   |
+LL | pub struct Layout {
+   | ^^^^^^^^^^^^^^^^^
+note: struct `Layout` is defined in the current crate
+  --> $DIR/alloc-error-handler-bad-signature-2.rs:7:1
+   |
+LL | struct Layout;
+   | ^^^^^^^^^^^^^
+note: function defined here
+  --> $DIR/alloc-error-handler-bad-signature-2.rs:10:4
+   |
+LL | fn oom(
+   |    ^^^
 LL |     info: Layout,
-   |           ^^^^^^
+   |     ------------
+   = note: this error originates in the attribute macro `alloc_error_handler` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error[E0308]: mismatched types
+  --> $DIR/alloc-error-handler-bad-signature-2.rs:10:1
+   |
+LL |    #[alloc_error_handler]
+   |    ---------------------- in this procedural macro expansion
+LL |    fn oom(
+   |   _^
+   |  |_|
+   | ||
+LL | ||     info: Layout,
+LL | || ) {
+   | || ^
+   | ||_|
+   | |  expected `!`, found `()`
+LL | |      loop {}
+LL | |  }
+   | |__- expected `!` because of return type
+   |
+   = note:   expected type `!`
+           found unit type `()`
+   = note: this error originates in the attribute macro `alloc_error_handler` (in Nightly builds, run with -Z macro-backtrace for more info)
 
 error: aborting due to 2 previous errors
 
+For more information about this error, try `rustc --explain E0308`.

src/test/ui/alloc-error/alloc-error-handler-bad-signature-3.rs

@@ -7,7 +7,7 @@
 struct Layout;
 
 #[alloc_error_handler]
-fn oom() -> ! { //~ ERROR function should have one argument
+fn oom() -> ! { //~ ERROR this function takes 0 arguments but 1 argument was supplied
     loop {}
 }
 

src/test/ui/alloc-error/alloc-error-handler-bad-signature-3.stderr

@@ -1,8 +1,25 @@
-error: function should have one argument
+error[E0061]: this function takes 0 arguments but 1 argument was supplied
   --> $DIR/alloc-error-handler-bad-signature-3.rs:10:1
    |
+LL |   #[alloc_error_handler]
+   |   ---------------------- in this procedural macro expansion
+LL |   fn oom() -> ! {
+   |  _-^^^^^^^^^^^^
+LL | |     loop {}
+LL | | }
+   | |_- argument of type `core::alloc::Layout` unexpected
+   |
+note: function defined here
+  --> $DIR/alloc-error-handler-bad-signature-3.rs:10:4
+   |
 LL | fn oom() -> ! {
-   | ^^^^^^^^^^^^^
+   |    ^^^
+   = note: this error originates in the attribute macro `alloc_error_handler` (in Nightly builds, run with -Z macro-backtrace for more info)
+help: remove the extra argument
+   |
+LL | fn oom() -> !() {
+   |              ++
 
 error: aborting due to previous error
 
+For more information about this error, try `rustc --explain E0061`.

src/test/ui/feature-gates/feature-gate-alloc-error-handler.rs

@@ -5,10 +5,12 @@
 
 use core::alloc::Layout;
 
-#[alloc_error_handler] //~ ERROR the `#[alloc_error_handler]` attribute is an experimental feature
+#[alloc_error_handler] //~ ERROR use of unstable library feature 'alloc_error_handler'
 fn oom(info: Layout) -> ! {
     loop {}
 }
 
 #[panic_handler]
-fn panic(_: &core::panic::PanicInfo) -> ! { loop {} }
+fn panic(_: &core::panic::PanicInfo) -> ! {
+    loop {}
+}

src/test/ui/feature-gates/feature-gate-alloc-error-handler.stderr

@@ -1,8 +1,8 @@
-error[E0658]: the `#[alloc_error_handler]` attribute is an experimental feature
+error[E0658]: use of unstable library feature 'alloc_error_handler'
-  --> $DIR/feature-gate-alloc-error-handler.rs:8:1
+  --> $DIR/feature-gate-alloc-error-handler.rs:8:3
    |
 LL | #[alloc_error_handler]
-   | ^^^^^^^^^^^^^^^^^^^^^^
+   |   ^^^^^^^^^^^^^^^^^^^
    |
    = note: see issue #51540 <https://github.com/rust-lang/rust/issues/51540> for more information
    = help: add `#![feature(alloc_error_handler)]` to the crate attributes to enable