From 778330b242fef436b4d1bc95e6fb648788092aeb Mon Sep 17 00:00:00 2001
From: GnomedDev <david2005thomas@gmail.com>
Date: Sat, 6 Apr 2024 18:57:46 +0100
Subject: [PATCH] Add UI test to ensure no more allocations pre-main

---
 tests/ui/runtime/aborting-alloc.rs            | 26 +++++++++++++++++++
 tests/ui/runtime/no-allocation-before-main.rs | 13 ++++++++++
 2 files changed, 39 insertions(+)
 create mode 100644 tests/ui/runtime/aborting-alloc.rs
 create mode 100644 tests/ui/runtime/no-allocation-before-main.rs

diff --git a/tests/ui/runtime/aborting-alloc.rs b/tests/ui/runtime/aborting-alloc.rs
new file mode 100644
index 0000000000000..f078e7dc3f60e
--- /dev/null
+++ b/tests/ui/runtime/aborting-alloc.rs
@@ -0,0 +1,26 @@
+//! Helper for 'no-allocation-before-main'.
+//!
+//! This also contains a meta-test to make sure that the AbortingAllocator does indeed abort.
+//!
+//! -Cprefer-dynamic=no is required as otherwise #[global_allocator] does nothing.
+//@ run-fail
+//@ compile-flags: -Cprefer-dynamic=no
+
+pub struct AbortingAllocator;
+
+unsafe impl std::alloc::GlobalAlloc for AbortingAllocator {
+    unsafe fn alloc(&self, _: std::alloc::Layout) -> *mut u8 {
+        std::process::abort()
+    }
+
+    unsafe fn dealloc(&self, _: *mut u8, _: std::alloc::Layout) {
+        std::process::abort()
+    }
+}
+
+#[global_allocator]
+static ALLOCATOR: AbortingAllocator = AbortingAllocator;
+
+fn main() {
+    std::hint::black_box(String::from("An allocation"));
+}
diff --git a/tests/ui/runtime/no-allocation-before-main.rs b/tests/ui/runtime/no-allocation-before-main.rs
new file mode 100644
index 0000000000000..6c4e06d946fcb
--- /dev/null
+++ b/tests/ui/runtime/no-allocation-before-main.rs
@@ -0,0 +1,13 @@
+//! Tests that a program with no body does not allocate.
+//!
+//! The initial runtime should not allocate for performance/binary size reasons.
+//!
+//! -Cprefer-dynamic=no is required as otherwise #[global_allocator] does nothing.
+//@ run-pass
+//@ compile-flags: -Cprefer-dynamic=no
+
+#[allow(dead_code)]
+#[path = "aborting-alloc.rs"]
+mod aux;
+
+fn main() {}