Auto merge of #3480 - RalfJung:alloc_error_handler, r=RalfJung

directly call handle_alloc_error

Also test more codepaths. There's like 5 different things that can happen on allocation failure! Between `-Zoom`, `#[alloc_error_handler]`, and `set_alloc_error_hook`, we have 3 layers of behavior overrides. It's all a bit messy.

https://github.com/rust-lang/rust/pull/112331 seems intended to clean this up, but has not yet reached consensus.

16 files changed, 171 insertions, 120 deletions

diff --git a/src/tools/miri/tests/fail/alloc/alloc_error_handler.rs b/src/tools/miri/tests/fail/alloc/alloc_error_handler.rs
index dc8e8c73800..2097126e16b 100644
--- a/src/tools/miri/tests/fail/alloc/alloc_error_handler.rs
+++ b/src/tools/miri/tests/fail/alloc/alloc_error_handler.rs
@@ -4,22 +4,7 @@
 #![feature(allocator_api)]
 
 use std::alloc::*;
-use std::ptr::NonNull;
-
-struct BadAlloc;
-
-// Create a failing allocator; Miri's native allocator never fails so this is the only way to
-// actually call the alloc error handler.
-unsafe impl Allocator for BadAlloc {
-    fn allocate(&self, _l: Layout) -> Result<NonNull<[u8]>, AllocError> {
-        Err(AllocError)
-    }
-
-    unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: Layout) {
-        unreachable!();
-    }
-}
 
 fn main() {
-    let _b = Box::new_in(0, BadAlloc);
+    handle_alloc_error(Layout::for_value(&0));
 }
diff --git a/src/tools/miri/tests/fail/alloc/alloc_error_handler.stderr b/src/tools/miri/tests/fail/alloc/alloc_error_handler.stderr
index f9d8e80c0fa..d1731a0f420 100644
--- a/src/tools/miri/tests/fail/alloc/alloc_error_handler.stderr
+++ b/src/tools/miri/tests/fail/alloc/alloc_error_handler.stderr
@@ -12,12 +12,10 @@ LL |     ABORT();
    = note: inside `std::alloc::_::__rg_oom` at RUSTLIB/std/src/alloc.rs:LL:CC
    = note: inside `std::alloc::handle_alloc_error::rt_error` at RUSTLIB/alloc/src/alloc.rs:LL:CC
    = note: inside `std::alloc::handle_alloc_error` at RUSTLIB/alloc/src/alloc.rs:LL:CC
-   = note: inside `std::boxed::Box::<i32, BadAlloc>::new_uninit_in` at RUSTLIB/alloc/src/boxed.rs:LL:CC
-   = note: inside `std::boxed::Box::<i32, BadAlloc>::new_in` at RUSTLIB/alloc/src/boxed.rs:LL:CC
 note: inside `main`
   --> $DIR/alloc_error_handler.rs:LL:CC
    |
-LL |     let _b = Box::new_in(0, BadAlloc);
+LL |     handle_alloc_error(Layout::for_value(&0));
    | ^
 
 note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
diff --git a/src/tools/miri/tests/fail/alloc/alloc_error_handler_custom.rs b/src/tools/miri/tests/fail/alloc/alloc_error_handler_custom.rs
new file mode 100644
index 00000000000..babdb73f093
--- /dev/null
+++ b/src/tools/miri/tests/fail/alloc/alloc_error_handler_custom.rs
@@ -0,0 +1,49 @@
+//@compile-flags: -Cpanic=abort
+#![feature(start, core_intrinsics)]
+#![feature(alloc_error_handler)]
+#![feature(allocator_api)]
+#![no_std]
+
+extern crate alloc;
+
+use alloc::alloc::*;
+use core::fmt::Write;
+
+#[path = "../../utils/mod.no_std.rs"]
+mod utils;
+
+#[alloc_error_handler]
+fn alloc_error_handler(layout: Layout) -> ! {
+    let _ = writeln!(utils::MiriStderr, "custom alloc error handler: {layout:?}");
+    core::intrinsics::abort(); //~ERROR: aborted
+}
+
+// rustc requires us to provide some more things that aren't actually used by this test
+mod plumbing {
+    use super::*;
+
+    #[panic_handler]
+    fn panic_handler(_: &core::panic::PanicInfo) -> ! {
+        loop {}
+    }
+
+    struct NoAlloc;
+
+    unsafe impl GlobalAlloc for NoAlloc {
+        unsafe fn alloc(&self, _: Layout) -> *mut u8 {
+            unreachable!();
+        }
+
+        unsafe fn dealloc(&self, _: *mut u8, _: Layout) {
+            unreachable!();
+        }
+    }
+
+    #[global_allocator]
+    static GLOBAL: NoAlloc = NoAlloc;
+}
+
+#[start]
+fn start(_: isize, _: *const *const u8) -> isize {
+    handle_alloc_error(Layout::for_value(&0));
+}
diff --git a/src/tools/miri/tests/fail/alloc/alloc_error_handler_custom.stderr b/src/tools/miri/tests/fail/alloc/alloc_error_handler_custom.stderr
new file mode 100644
index 00000000000..5d9c2e2fb4c
--- /dev/null
+++ b/src/tools/miri/tests/fail/alloc/alloc_error_handler_custom.stderr
@@ -0,0 +1,27 @@
+custom alloc error handler: Layout { size: 4, align: 4 (1 << 2) }
+error: abnormal termination: the program aborted execution
+  --> $DIR/alloc_error_handler_custom.rs:LL:CC
+   |
+LL |     core::intrinsics::abort();
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^ the program aborted execution
+   |
+   = note: BACKTRACE:
+   = note: inside `alloc_error_handler` at $DIR/alloc_error_handler_custom.rs:LL:CC
+note: inside `_::__rg_oom`
+  --> $DIR/alloc_error_handler_custom.rs:LL:CC
+   |
+LL | #[alloc_error_handler]
+   | ---------------------- in this procedural macro expansion
+LL | fn alloc_error_handler(layout: Layout) -> ! {
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   = note: inside `alloc::alloc::handle_alloc_error::rt_error` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `alloc::alloc::handle_alloc_error` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+note: inside `start`
+  --> $DIR/alloc_error_handler_custom.rs:LL:CC
+   |
+LL |     handle_alloc_error(Layout::for_value(&0));
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   = 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 1 previous error
+
diff --git a/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.rs b/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.rs
index 8103296f47b..18a8a61f22f 100644
--- a/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.rs
+++ b/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.rs
@@ -7,29 +7,17 @@
 extern crate alloc;
 
 use alloc::alloc::*;
-use alloc::boxed::Box;
-use core::ptr::NonNull;
+use core::fmt::Write;
 
-struct BadAlloc;
+#[path = "../../utils/mod.no_std.rs"]
+mod utils;
 
-// Create a failing allocator; that is the only way to actually call the alloc error handler.
-unsafe impl Allocator for BadAlloc {
-    fn allocate(&self, _l: Layout) -> Result<NonNull<[u8]>, AllocError> {
-        Err(AllocError)
-    }
-
-    unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: Layout) {
-        unreachable!();
-    }
-}
+// The default no_std alloc_error_handler is a panic.
 
-#[alloc_error_handler]
-fn alloc_error_handler(_: Layout) -> ! {
-    extern "Rust" {
-        fn miri_write_to_stderr(bytes: &[u8]);
-    }
-    let msg = "custom alloc error handler called!\n";
-    unsafe { miri_write_to_stderr(msg.as_bytes()) };
+#[panic_handler]
+fn panic_handler(panic_info: &core::panic::PanicInfo) -> ! {
+    let _ = writeln!(utils::MiriStderr, "custom panic handler called!");
+    let _ = writeln!(utils::MiriStderr, "{panic_info}");
     core::intrinsics::abort(); //~ERROR: aborted
 }
 
@@ -37,11 +25,6 @@ fn alloc_error_handler(_: Layout) -> ! {
 mod plumbing {
     use super::*;
 
-    #[panic_handler]
-    fn panic_handler(_: &core::panic::PanicInfo) -> ! {
-        core::intrinsics::abort();
-    }
-
     struct NoAlloc;
 
     unsafe impl GlobalAlloc for NoAlloc {
@@ -60,6 +43,5 @@ mod plumbing {
 
 #[start]
 fn start(_: isize, _: *const *const u8) -> isize {
-    let _b = Box::new_in(0, BadAlloc);
-    0
+    handle_alloc_error(Layout::for_value(&0));
 }
diff --git a/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.stderr b/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.stderr
index b40ffb70121..6b98f6f75d8 100644
--- a/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.stderr
+++ b/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.stderr
@@ -1,4 +1,6 @@
-custom alloc error handler called!
+custom panic handler called!
+panicked at RUSTLIB/alloc/src/alloc.rs:LL:CC:
+memory allocation of 4 bytes failed
 error: abnormal termination: the program aborted execution
   --> $DIR/alloc_error_handler_no_std.rs:LL:CC
    |
@@ -6,24 +8,17 @@ LL |     core::intrinsics::abort();
    |     ^^^^^^^^^^^^^^^^^^^^^^^^^ the program aborted execution
    |
    = note: BACKTRACE:
-   = note: inside `alloc_error_handler` at $DIR/alloc_error_handler_no_std.rs:LL:CC
-note: inside `_::__rg_oom`
-  --> $DIR/alloc_error_handler_no_std.rs:LL:CC
-   |
-LL | #[alloc_error_handler]
-   | ---------------------- in this procedural macro expansion
-LL | fn alloc_error_handler(_: Layout) -> ! {
-   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   = note: inside `panic_handler` at $DIR/alloc_error_handler_no_std.rs:LL:CC
+   = note: inside `alloc::alloc::__alloc_error_handler::__rdl_oom` at RUSTLIB/alloc/src/alloc.rs:LL:CC
    = note: inside `alloc::alloc::handle_alloc_error::rt_error` at RUSTLIB/alloc/src/alloc.rs:LL:CC
    = note: inside `alloc::alloc::handle_alloc_error` at RUSTLIB/alloc/src/alloc.rs:LL:CC
-   = note: inside `alloc::boxed::Box::<i32, BadAlloc>::new_uninit_in` at RUSTLIB/alloc/src/boxed.rs:LL:CC
-   = note: inside `alloc::boxed::Box::<i32, BadAlloc>::new_in` at RUSTLIB/alloc/src/boxed.rs:LL:CC
 note: inside `start`
   --> $DIR/alloc_error_handler_no_std.rs:LL:CC
    |
-LL |     let _b = Box::new_in(0, BadAlloc);
-   |              ^^^^^^^^^^^^^^^^^^^^^^^^
-   = note: this error originates in the attribute macro `alloc_error_handler` (in Nightly builds, run with -Z macro-backtrace for more info)
+LL |     handle_alloc_error(Layout::for_value(&0));
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
 
 error: aborting due to 1 previous error
 
diff --git a/src/tools/miri/tests/fail/panic/no_std.rs b/src/tools/miri/tests/fail/panic/no_std.rs
index 26cc0b27821..4d32b6d7461 100644
--- a/src/tools/miri/tests/fail/panic/no_std.rs
+++ b/src/tools/miri/tests/fail/panic/no_std.rs
@@ -1,27 +1,11 @@
+//@compile-flags: -Cpanic=abort
 #![feature(start, core_intrinsics)]
 #![no_std]
-//@compile-flags: -Cpanic=abort
-
-// Plumbing to let us use `writeln!` to host stderr:
-
-extern "Rust" {
-    fn miri_write_to_stderr(bytes: &[u8]);
-}
-
-struct HostErr;
 
 use core::fmt::Write;
 
-impl Write for HostErr {
-    fn write_str(&mut self, s: &str) -> core::fmt::Result {
-        unsafe {
-            miri_write_to_stderr(s.as_bytes());
-        }
-        Ok(())
-    }
-}
-
-// Aaaand the test:
+#[path = "../../utils/mod.no_std.rs"]
+mod utils;
 
 #[start]
 fn start(_: isize, _: *const *const u8) -> isize {
@@ -30,6 +14,6 @@ fn start(_: isize, _: *const *const u8) -> isize {
 
 #[panic_handler]
 fn panic_handler(panic_info: &core::panic::PanicInfo) -> ! {
-    writeln!(HostErr, "{panic_info}").ok();
+    writeln!(utils::MiriStderr, "{panic_info}").ok();
     core::intrinsics::abort(); //~ ERROR: the program aborted execution
 }
diff --git a/src/tools/miri/tests/panic/alloc_error_handler_hook.rs b/src/tools/miri/tests/panic/alloc_error_handler_hook.rs
new file mode 100644
index 00000000000..a1eadb45fd1
--- /dev/null
+++ b/src/tools/miri/tests/panic/alloc_error_handler_hook.rs
@@ -0,0 +1,20 @@
+#![feature(allocator_api, alloc_error_hook)]
+
+use std::alloc::*;
+
+struct Bomb;
+impl Drop for Bomb {
+    fn drop(&mut self) {
+        eprintln!("yes we are unwinding!");
+    }
+}
+
+#[allow(unreachable_code, unused_variables)]
+fn main() {
+    // This is a particularly tricky hook, since it unwinds, which the default one does not.
+    set_alloc_error_hook(|_layout| panic!("alloc error hook called"));
+
+    let bomb = Bomb;
+    handle_alloc_error(Layout::for_value(&0));
+    std::mem::forget(bomb); // defuse unwinding bomb
+}
diff --git a/src/tools/miri/tests/panic/alloc_error_handler_hook.stderr b/src/tools/miri/tests/panic/alloc_error_handler_hook.stderr
new file mode 100644
index 00000000000..01c286fe3d2
--- /dev/null
+++ b/src/tools/miri/tests/panic/alloc_error_handler_hook.stderr
@@ -0,0 +1,4 @@
+thread 'main' panicked at $DIR/alloc_error_handler_hook.rs:LL:CC:
+alloc error hook called
+note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
+yes we are unwinding!
diff --git a/src/tools/miri/tests/panic/alloc_error_handler_panic.rs b/src/tools/miri/tests/panic/alloc_error_handler_panic.rs
index 186c9667a97..c434e8d3227 100644
--- a/src/tools/miri/tests/panic/alloc_error_handler_panic.rs
+++ b/src/tools/miri/tests/panic/alloc_error_handler_panic.rs
@@ -2,21 +2,6 @@
 #![feature(allocator_api)]
 
 use std::alloc::*;
-use std::ptr::NonNull;
-
-struct BadAlloc;
-
-// Create a failing allocator; Miri's native allocator never fails so this is the only way to
-// actually call the alloc error handler.
-unsafe impl Allocator for BadAlloc {
-    fn allocate(&self, _l: Layout) -> Result<NonNull<[u8]>, AllocError> {
-        Err(AllocError)
-    }
-
-    unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: Layout) {
-        unreachable!();
-    }
-}
 
 struct Bomb;
 impl Drop for Bomb {
@@ -25,8 +10,9 @@ impl Drop for Bomb {
     }
 }
 
+#[allow(unreachable_code, unused_variables)]
 fn main() {
     let bomb = Bomb;
-    let _b = Box::new_in(0, BadAlloc);
+    handle_alloc_error(Layout::for_value(&0));
     std::mem::forget(bomb); // defuse unwinding bomb
 }
diff --git a/src/tools/miri/tests/pass/no_std.rs b/src/tools/miri/tests/pass/no_std.rs
index 3c98ee50aa9..fc1c16f5fb9 100644
--- a/src/tools/miri/tests/pass/no_std.rs
+++ b/src/tools/miri/tests/pass/no_std.rs
@@ -2,30 +2,14 @@
 #![feature(start)]
 #![no_std]
 
-// Plumbing to let us use `writeln!` to host stdout:
-
-extern "Rust" {
-    fn miri_write_to_stdout(bytes: &[u8]);
-}
-
-struct Host;
-
 use core::fmt::Write;
 
-impl Write for Host {
-    fn write_str(&mut self, s: &str) -> core::fmt::Result {
-        unsafe {
-            miri_write_to_stdout(s.as_bytes());
-        }
-        Ok(())
-    }
-}
-
-// Aaaand the test:
+#[path = "../utils/mod.no_std.rs"]
+mod utils;
 
 #[start]
 fn start(_: isize, _: *const *const u8) -> isize {
-    writeln!(Host, "hello, world!").unwrap();
+    writeln!(utils::MiriStdout, "hello, world!").unwrap();
     0
 }
 
diff --git a/src/tools/miri/tests/pass/tree_borrows/reserved.rs b/src/tools/miri/tests/pass/tree_borrows/reserved.rs
index 87ce91a809f..f93cac8361e 100644
--- a/src/tools/miri/tests/pass/tree_borrows/reserved.rs
+++ b/src/tools/miri/tests/pass/tree_borrows/reserved.rs
@@ -27,9 +27,8 @@ fn main() {
     }
 }
 
-unsafe fn print(msg: &str) {
-    utils::miri_write_to_stderr(msg.as_bytes());
-    utils::miri_write_to_stderr("\n".as_bytes());
+fn print(msg: &str) {
+    eprintln!("{msg}");
 }
 
 unsafe fn read_second<T>(x: &mut T, y: *mut u8) {
diff --git a/src/tools/miri/tests/utils/io.rs b/src/tools/miri/tests/utils/io.rs
new file mode 100644
index 00000000000..e3eaa6c468a
--- /dev/null
+++ b/src/tools/miri/tests/utils/io.rs
@@ -0,0 +1,25 @@
+use core::fmt::{self, Write};
+
+use super::miri_extern;
+
+pub struct MiriStderr;
+
+impl Write for MiriStderr {
+    fn write_str(&mut self, s: &str) -> fmt::Result {
+        unsafe {
+            miri_extern::miri_write_to_stderr(s.as_bytes());
+        }
+        Ok(())
+    }
+}
+
+pub struct MiriStdout;
+
+impl Write for MiriStdout {
+    fn write_str(&mut self, s: &str) -> fmt::Result {
+        unsafe {
+            miri_extern::miri_write_to_stdout(s.as_bytes());
+        }
+        Ok(())
+    }
+}
diff --git a/src/tools/miri/tests/utils/miri_extern.rs b/src/tools/miri/tests/utils/miri_extern.rs
index e2983f6c71a..d6c43b18821 100644
--- a/src/tools/miri/tests/utils/miri_extern.rs
+++ b/src/tools/miri/tests/utils/miri_extern.rs
@@ -133,8 +133,8 @@ extern "Rust" {
     /// with a null terminator.
     /// Returns 0 if the `out` buffer was large enough, and the required size otherwise.
     pub fn miri_host_to_target_path(
-        path: *const std::ffi::c_char,
-        out: *mut std::ffi::c_char,
+        path: *const core::ffi::c_char,
+        out: *mut core::ffi::c_char,
         out_size: usize,
     ) -> usize;
 
diff --git a/src/tools/miri/tests/utils/mod.no_std.rs b/src/tools/miri/tests/utils/mod.no_std.rs
new file mode 100644
index 00000000000..aaf2bf50c4e
--- /dev/null
+++ b/src/tools/miri/tests/utils/mod.no_std.rs
@@ -0,0 +1,11 @@
+#![allow(dead_code)]
+#![allow(unused_imports)]
+
+#[macro_use]
+mod macros;
+
+mod io;
+mod miri_extern;
+
+pub use self::io::*;
+pub use self::miri_extern::*;
diff --git a/src/tools/miri/tests/utils/mod.rs b/src/tools/miri/tests/utils/mod.rs
index cb9380f5753..138ada4e20d 100644
--- a/src/tools/miri/tests/utils/mod.rs
+++ b/src/tools/miri/tests/utils/mod.rs
@@ -5,9 +5,11 @@
 mod macros;
 
 mod fs;
+mod io;
 mod miri_extern;
 
 pub use self::fs::*;
+pub use self::io::*;
 pub use self::miri_extern::*;
 
 pub fn run_provenance_gc() {