Revert "Remove #[alloc_error_handler] from the compiler and library"

This reverts commit abc0660.

Author: matthiaskrgr

alloc/src/alloc.rs

@@ -397,7 +397,7 @@ fn rust_oom(layout: Layout) -> ! {
         #[lang = "panic_impl"]
         fn panic_impl(pi: &core::panic::PanicInfo<'_>) -> !;
 
-        // This symbol is emitted by rustc .
+        // This symbol is emitted by rustc next to __rust_alloc_error_handler.
         // Its value depends on the -Zoom={panic,abort} compiler option.
         static __rust_alloc_error_handler_should_panic: u8;
     }
@@ -426,6 +426,13 @@ fn rust_oom(layout: Layout) -> ! {
 /// Callers of memory allocation APIs wishing to abort computation
 /// in response to an allocation error are encouraged to call this function,
 /// rather than directly invoking `panic!` or similar.
+///
+/// The default behavior of this function is to print a message to standard error
+/// and abort the process.
+/// It can be replaced with [`set_alloc_error_hook`] and [`take_alloc_error_hook`].
+///
+/// [`set_alloc_error_hook`]: ../../std/alloc/fn.set_alloc_error_hook.html
+/// [`take_alloc_error_hook`]: ../../std/alloc/fn.take_alloc_error_hook.html
 #[stable(feature = "global_alloc", since = "1.28.0")]
 #[rustc_const_unstable(feature = "const_alloc_error", issue = "92523")]
 #[cfg(all(not(no_global_oom_handling), not(test)))]
@@ -446,7 +453,6 @@ pub const fn handle_alloc_error(layout: Layout) -> ! {
 #[cfg(all(not(no_global_oom_handling), test))]
 pub use std::alloc::handle_alloc_error;
 
-#[cfg(bootstrap)]
 #[cfg(all(not(no_global_oom_handling), not(test)))]
 #[doc(hidden)]
 #[allow(unused_attributes)]

core/src/macros/mod.rs

@@ -1531,6 +1531,16 @@ 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).
+    #[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",

core/src/prelude/v1.rs

@@ -76,7 +76,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).
 #[stable(feature = "builtin_macro_prelude", since = "1.38.0")]
-pub use crate::macros::builtin::{bench, derive, global_allocator, test, test_case};
+pub use crate::macros::builtin::{
+    alloc_error_handler, bench, derive, global_allocator, test, test_case,
+};
 
 #[unstable(feature = "derive_const", issue = "none")]
 pub use crate::macros::builtin::derive_const;

std/src/alloc.rs

@@ -57,8 +57,9 @@
 #![stable(feature = "alloc_module", since = "1.28.0")]
 
 use core::intrinsics;
-use core::ptr;
 use core::ptr::NonNull;
+use core::sync::atomic::{AtomicPtr, Ordering};
+use core::{mem, ptr};
 
 #[stable(feature = "alloc_module", since = "1.28.0")]
 #[doc(inline)]
@@ -285,6 +286,76 @@ unsafe impl Allocator for System {
     }
 }
 
+static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
+
+/// Registers a custom allocation error hook, replacing any that was previously registered.
+///
+/// The allocation error hook is invoked when an infallible memory allocation fails, before
+/// the runtime aborts. The default hook prints a message to standard error,
+/// but this behavior can be customized with the [`set_alloc_error_hook`] and
+/// [`take_alloc_error_hook`] functions.
+///
+/// The hook is provided with a `Layout` struct which contains information
+/// about the allocation that failed.
+///
+/// The allocation error hook is a global resource.
+///
+/// # Examples
+///
+/// ```
+/// #![feature(alloc_error_hook)]
+///
+/// use std::alloc::{Layout, set_alloc_error_hook};
+///
+/// fn custom_alloc_error_hook(layout: Layout) {
+///    panic!("memory allocation of {} bytes failed", layout.size());
+/// }
+///
+/// set_alloc_error_hook(custom_alloc_error_hook);
+/// ```
+#[unstable(feature = "alloc_error_hook", issue = "51245")]
+pub fn set_alloc_error_hook(hook: fn(Layout)) {
+    HOOK.store(hook as *mut (), Ordering::SeqCst);
+}
+
+/// Unregisters the current allocation error hook, returning it.
+///
+/// *See also the function [`set_alloc_error_hook`].*
+///
+/// If no custom hook is registered, the default hook will be returned.
+#[unstable(feature = "alloc_error_hook", issue = "51245")]
+pub fn take_alloc_error_hook() -> fn(Layout) {
+    let hook = HOOK.swap(ptr::null_mut(), Ordering::SeqCst);
+    if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } }
+}
+
+fn default_alloc_error_hook(layout: Layout) {
+    extern "Rust" {
+        // This symbol is emitted by rustc next to __rust_alloc_error_handler.
+        // Its value depends on the -Zoom={panic,abort} compiler option.
+        static __rust_alloc_error_handler_should_panic: u8;
+    }
+
+    #[allow(unused_unsafe)]
+    if unsafe { __rust_alloc_error_handler_should_panic != 0 } {
+        panic!("memory allocation of {} bytes failed", layout.size());
+    } else {
+        rtprintpanic!("memory allocation of {} bytes failed\n", layout.size());
+    }
+}
+
+#[cfg(not(test))]
+#[doc(hidden)]
+#[alloc_error_handler]
+#[unstable(feature = "alloc_internals", issue = "none")]
+pub fn rust_oom(layout: Layout) -> ! {
+    let hook = HOOK.load(Ordering::SeqCst);
+    let hook: fn(Layout) =
+        if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } };
+    hook(layout);
+    crate::process::abort()
+}
+
 #[cfg(not(test))]
 #[doc(hidden)]
 #[allow(unused_attributes)]

std/src/lib.rs

@@ -236,6 +236,7 @@
 //
 // Language features:
 // tidy-alphabetical-start
+#![feature(alloc_error_handler)]
 #![feature(allocator_internals)]
 #![feature(allow_internal_unsafe)]
 #![feature(allow_internal_unstable)]

std/src/prelude/v1.rs

@@ -60,7 +60,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).
 #[stable(feature = "builtin_macro_prelude", since = "1.38.0")]
-pub use core::prelude::v1::{bench, derive, global_allocator, test, test_case};
+pub use core::prelude::v1::{
+    alloc_error_handler, bench, derive, global_allocator, test, test_case,
+};
 
 #[unstable(feature = "derive_const", issue = "none")]
 pub use core::prelude::v1::derive_const;