Auto merge of #157153 - joboet:global_allocator, r=<try>

allow `Allocator`s to be used as `#[global_allocator]`s


try-job: x86_64-msvc-1

Author: bors

library/core/src/alloc/global.rs

@@ -1,4 +1,7 @@
+use super::{AllocError, GlobalAllocator};
 use crate::alloc::Layout;
+use crate::hint::assert_unchecked;
+use crate::ptr::NonNull;
 use crate::{cmp, ptr};
 
 /// A memory allocator that can be registered as the standard library’s default
@@ -301,3 +304,72 @@ pub unsafe trait GlobalAlloc {
         new_ptr
     }
 }
+
+/// Allows all [`GlobalAllocator`]s to be used with the legacy [`GlobalAlloc`] interface.
+#[stable(feature = "global_alloc", since = "1.28.0")]
+unsafe impl<A> GlobalAlloc for A
+where
+    A: GlobalAllocator + ?Sized,
+{
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        // SAFETY: guaranteed by the caller.
+        // This might lead to the removal of zero-size checks inside the
+        // `Allocator` implementation.
+        unsafe { assert_unchecked(layout.size() != 0) };
+        match self.allocate(layout) {
+            Ok(ptr) => ptr.cast().as_ptr(),
+            Err(AllocError) => ptr::null_mut(),
+        }
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        // SAFETY: guaranteed by the caller.
+        unsafe { assert_unchecked(layout.size() != 0) };
+        // SAFETY: only non-null pointers can be currently allocated.
+        let ptr = unsafe { NonNull::new_unchecked(ptr) };
+        // SAFETY: guaranteed by caller.
+        unsafe { self.deallocate(ptr, layout) };
+    }
+
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        // SAFETY: guaranteed by the caller.
+        unsafe { assert_unchecked(layout.size() != 0) };
+        match self.allocate_zeroed(layout) {
+            Ok(ptr) => ptr.cast().as_ptr(),
+            Err(AllocError) => ptr::null_mut(),
+        }
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        // SAFETY: guaranteed by the caller.
+        unsafe { assert_unchecked(layout.size() != 0) };
+        // SAFETY: guaranteed by the caller.
+        unsafe { assert_unchecked(new_size != 0) };
+
+        // SAFETY: only non-null pointers can be currently allocated.
+        let ptr = unsafe { NonNull::new_unchecked(ptr) };
+        let alignment = layout.alignment();
+        // SAFETY: the caller must ensure that the `new_size` does not overflow
+        // when rounded up to the next multiple of `alignment`.
+        let new_layout = unsafe { Layout::from_size_alignment_unchecked(new_size, alignment) };
+
+        // SAFETY:
+        // Two preconditions are guaranteed by the caller:
+        // * `ptr` is currently allocated with this allocator.
+        // * `layout` fits the block of memory.
+        // The size precondition is upheld by selecting between `grow` and `shrink`
+        // based on the size.
+        let ptr = unsafe {
+            if new_size >= layout.size() {
+                self.grow(ptr, layout, new_layout)
+            } else {
+                self.shrink(ptr, layout, new_layout)
+            }
+        };
+
+        match ptr {
+            Ok(ptr) => ptr.cast().as_ptr(),
+            Err(AllocError) => ptr::null_mut(),
+        }
+    }
+}

library/core/src/alloc/mod.rs

@@ -447,6 +447,96 @@ pub const unsafe trait Allocator {
     }
 }
 
+/// An [`Allocator`] that can be registered as the standard library’s default
+/// through the `#[global_allocator]` attribute.
+///
+/// Types implementing this trait can be used as the default allocator for
+/// memory allocations through `Box`, `Vec` and the collection types. For
+/// instance, the `System` allocator implements this trait, and thus can be
+/// explicitly set as the default like so:
+/// ```
+/// use std::alloc::System;
+///
+/// #[global_allocator]
+/// static ALLOCATOR: System = System;
+/// ```
+///
+/// The `Global` allocator forwards all memory allocation requests to the
+/// `static` annotated with `#[global_allocator]`. Hence, `Global` does not
+/// implement `GlobalAllocator` itself, as that would lead to infinite recursion.
+///
+/// # Note to implementors
+///
+/// This trait is used to prevent the infinite recursion that would occur if the
+/// default allocator were to attempt to allocate memory through `Global` (and
+/// thus from itself).
+///
+/// When to implement this trait:
+/// * for custom global allocators that only use system memory allocation
+///   services.
+/// * for allocators that wrap another allocator that implements `GlobalAllocator`.
+///
+/// When **not** to implement this trait:
+/// * for wrappers of arbitrary allocators (which might end up being `Global`,
+///   leading to infinite recursion).
+///
+/// # Safety
+///
+/// In addition to the safety requirements of `Allocator`, global allocators are
+/// subject to some additional constraints:
+///
+/// * It's undefined behavior if global allocators unwind. This restriction may
+///   be lifted in the future, but currently a panic from any of these
+///   functions may lead to memory unsafety.
+///
+/// * You must not rely on allocations actually happening, even if there are explicit
+///   heap allocations in the source. The optimizer may detect unused allocations that it can either
+///   eliminate entirely or move to the stack and thus never invoke the allocator. The
+///   optimizer may further assume that allocation is infallible, so code that used to fail due
+///   to allocator failures may now suddenly work because the optimizer worked around the
+///   need for an allocation. More concretely, the following code example is unsound, irrespective
+///   of whether your custom allocator allows counting how many allocations have happened.
+///
+///   ```rust,ignore (unsound and has placeholders)
+///   drop(Box::new(42));
+///   let number_of_heap_allocs = /* call private allocator API */;
+///   unsafe { std::hint::assert_unchecked(number_of_heap_allocs > 0); }
+///   ```
+///
+///   Note that the optimizations mentioned above are not the only
+///   optimization that can be applied. You may generally not rely on heap allocations
+///   happening if they can be removed without changing program behavior.
+///   Whether allocations happen or not is not part of the program behavior, even if it
+///   could be detected via an allocator that tracks allocations by printing or otherwise
+///   having side effects.
+///
+/// # Re-entrance
+///
+/// When implementing a global allocator, one has to be careful not to create an infinitely recursive
+/// implementation by accident, as many constructs in the Rust standard library may allocate in
+/// their implementation. For example, on some platforms, [`std::sync::Mutex`] may allocate, so using
+/// it is highly problematic in a global allocator.
+///
+/// For this reason, one should generally stick to library features available through
+/// [`core`], and avoid using [`std`] in a global allocator. A few features from [`std`] are
+/// guaranteed to not use `#[global_allocator]` to allocate:
+///
+///  - [`std::thread_local`],
+///  - [`std::thread::current`],
+///  - [`std::thread::park`] and [`std::thread::Thread`]'s [`unpark`] method and
+/// [`Clone`] implementation.
+///
+/// [`std`]: ../../std/index.html
+/// [`std::sync::Mutex`]: ../../std/sync/struct.Mutex.html
+/// [`std::thread_local`]: ../../std/macro.thread_local.html
+/// [`std::thread::current`]: ../../std/thread/fn.current.html
+/// [`std::thread::park`]: ../../std/thread/fn.park.html
+/// [`std::thread::Thread`]: ../../std/thread/struct.Thread.html
+/// [`unpark`]: ../../std/thread/struct.Thread.html#method.unpark
+#[unstable(feature = "allocator_api", issue = "32838")]
+#[expect(multiple_supertrait_upcastable)]
+pub unsafe trait GlobalAllocator: Allocator + Sync + 'static {}
+
 #[unstable(feature = "allocator_api", issue = "32838")]
 #[rustc_const_unstable(feature = "const_heap", issue = "79597")]
 const unsafe impl<A> Allocator for &A

library/std/src/alloc.rs

@@ -63,14 +63,15 @@
 #![deny(unsafe_op_in_unsafe_fn)]
 #![stable(feature = "alloc_module", since = "1.28.0")]
 
-use core::ptr::NonNull;
-use core::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
-use core::{hint, mem, ptr};
-
 #[stable(feature = "alloc_module", since = "1.28.0")]
 #[doc(inline)]
 pub use alloc_crate::alloc::*;
 
+use crate::ptr::NonNull;
+use crate::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
+use crate::sys::alloc as imp;
+use crate::{hint, mem, ptr};
+
 /// The default memory allocator provided by the operating system.
 ///
 /// This is based on `malloc` on Unix platforms and `HeapAlloc` on Windows,
@@ -145,11 +146,7 @@ impl System {
             0 => Ok(layout.dangling_ptr().cast_slice(0)),
             // SAFETY: `layout` is non-zero in size,
             size => unsafe {
-                let raw_ptr = if zeroed {
-                    GlobalAlloc::alloc_zeroed(self, layout)
-                } else {
-                    GlobalAlloc::alloc(self, layout)
-                };
+                let raw_ptr = if zeroed { imp::alloc_zeroed(layout) } else { imp::alloc(layout) };
                 let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
                 Ok(ptr.cast_slice(size))
             },
@@ -182,7 +179,7 @@ impl System {
                 // `realloc` probably checks for `new_size >= old_layout.size()` or something similar.
                 hint::assert_unchecked(new_size >= old_layout.size());
 
-                let raw_ptr = GlobalAlloc::realloc(self, ptr.as_ptr(), old_layout, new_size);
+                let raw_ptr = imp::realloc(ptr.as_ptr(), old_layout, new_size);
                 let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
                 if zeroed {
                     raw_ptr.add(old_size).write_bytes(0, new_size - old_size);
@@ -205,8 +202,8 @@ impl System {
     }
 }
 
-// The Allocator impl checks the layout size to be non-zero and forwards to the GlobalAlloc impl,
-// which is in `std::sys::*::alloc`.
+// The Allocator impl checks the layout size to be non-zero and forwards to the
+// platform functions in `std::sys::*::alloc`.
 #[unstable(feature = "allocator_api", issue = "32838")]
 unsafe impl Allocator for System {
     #[inline]
@@ -224,7 +221,7 @@ unsafe impl Allocator for System {
         if layout.size() != 0 {
             // SAFETY: `layout` is non-zero in size,
             // other conditions must be upheld by the caller
-            unsafe { GlobalAlloc::dealloc(self, ptr.as_ptr(), layout) }
+            unsafe { imp::dealloc(ptr.as_ptr(), layout) }
         }
     }
 
@@ -274,7 +271,7 @@ unsafe impl Allocator for System {
                 // `realloc` probably checks for `new_size <= old_layout.size()` or something similar.
                 hint::assert_unchecked(new_size <= old_layout.size());
 
-                let raw_ptr = GlobalAlloc::realloc(self, ptr.as_ptr(), old_layout, new_size);
+                let raw_ptr = imp::realloc(ptr.as_ptr(), old_layout, new_size);
                 let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
                 Ok(ptr.cast_slice(new_size))
             },
@@ -294,6 +291,9 @@ unsafe impl Allocator for System {
     }
 }
 
+#[unstable(feature = "allocator_api", issue = "32838")]
+unsafe impl GlobalAllocator for System {}
+
 static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
 
 /// Registers a custom allocation error hook, replacing any that was previously registered.
@@ -435,7 +435,12 @@ pub fn rust_oom(layout: Layout) -> ! {
 #[allow(unused_attributes)]
 #[unstable(feature = "alloc_internals", issue = "none")]
 pub mod __default_lib_allocator {
-    use super::{GlobalAlloc, Layout, System};
+    use super::Layout;
+    // We call the system functions directly to avoid any overheads introduced
+    // by the roundtrip through `impl Allocator for System` and
+    // `impl<A: GlobalAllocator> GlobalAlloc for A`.
+    use crate::sys::alloc as imp;
+
     // These magic symbol names are used as a fallback for implementing the
     // `__rust_alloc` etc symbols (see `src/liballoc/alloc.rs`) when there is
     // no `#[global_allocator]` attribute.
@@ -452,15 +457,15 @@ pub mod __default_lib_allocator {
         // `GlobalAlloc::alloc`.
         unsafe {
             let layout = Layout::from_size_align_unchecked(size, align);
-            System.alloc(layout)
+            imp::alloc(layout)
         }
     }
 
     #[rustc_std_internal_symbol]
     pub unsafe extern "C" fn __rdl_dealloc(ptr: *mut u8, size: usize, align: usize) {
         // SAFETY: see the guarantees expected by `Layout::from_size_align` and
         // `GlobalAlloc::dealloc`.
-        unsafe { System.dealloc(ptr, Layout::from_size_align_unchecked(size, align)) }
+        unsafe { imp::dealloc(ptr, Layout::from_size_align_unchecked(size, align)) }
     }
 
     #[rustc_std_internal_symbol]
@@ -474,7 +479,7 @@ pub mod __default_lib_allocator {
         // `GlobalAlloc::realloc`.
         unsafe {
             let old_layout = Layout::from_size_align_unchecked(old_size, align);
-            System.realloc(ptr, old_layout, new_size)
+            imp::realloc(ptr, old_layout, new_size)
         }
     }
 
@@ -484,7 +489,7 @@ pub mod __default_lib_allocator {
         // `GlobalAlloc::alloc_zeroed`.
         unsafe {
             let layout = Layout::from_size_align_unchecked(size, align);
-            System.alloc_zeroed(layout)
+            imp::alloc_zeroed(layout)
         }
     }
 }

library/std/src/sys/alloc/hermit.rs

@@ -1,27 +1,24 @@
-use crate::alloc::{GlobalAlloc, Layout, System};
+use crate::alloc::Layout;
 
-#[stable(feature = "alloc_system_type", since = "1.28.0")]
-unsafe impl GlobalAlloc for System {
-    #[inline]
-    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        let size = layout.size();
-        let align = layout.align();
-        unsafe { hermit_abi::malloc(size, align) }
-    }
+#[inline]
+pub unsafe fn alloc(layout: Layout) -> *mut u8 {
+    let size = layout.size();
+    let align = layout.align();
+    unsafe { hermit_abi::malloc(size, align) }
+}
 
-    #[inline]
-    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-        let size = layout.size();
-        let align = layout.align();
-        unsafe {
-            hermit_abi::free(ptr, size, align);
-        }
+#[inline]
+pub unsafe fn dealloc(ptr: *mut u8, layout: Layout) {
+    let size = layout.size();
+    let align = layout.align();
+    unsafe {
+        hermit_abi::free(ptr, size, align);
     }
+}
 
-    #[inline]
-    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
-        let size = layout.size();
-        let align = layout.align();
-        unsafe { hermit_abi::realloc(ptr, size, align, new_size) }
-    }
+#[inline]
+pub unsafe fn realloc(ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+    let size = layout.size();
+    let align = layout.align();
+    unsafe { hermit_abi::realloc(ptr, size, align, new_size) }
 }