library/core/src/alloc/global.rs

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
/// through the `#[global_allocator]` attribute.
///
/// Some of the methods require that a memory block be *currently
/// allocated* via an allocator. This means that:
///
/// * the starting address for that memory block was previously
///   returned by a previous call to an allocation method
///   such as `alloc`, and
///
/// * the memory block has not been subsequently deallocated, where
///   blocks are deallocated either by being passed to a deallocation
///   method such as `dealloc` or by being
///   passed to a reallocation method that returns a non-null pointer.
///
/// # Example
///
/// ```standalone_crate
/// use std::alloc::{GlobalAlloc, Layout};
/// use std::cell::UnsafeCell;
/// use std::ptr::null_mut;
/// use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
///
/// const ARENA_SIZE: usize = 128 * 1024;
/// const MAX_SUPPORTED_ALIGN: usize = 4096;
/// #[repr(C, align(4096))] // 4096 == MAX_SUPPORTED_ALIGN
/// struct SimpleAllocator {
///     arena: UnsafeCell<[u8; ARENA_SIZE]>,
///     remaining: AtomicUsize, // we allocate from the top, counting down
/// }
///
/// #[global_allocator]
/// static ALLOCATOR: SimpleAllocator = SimpleAllocator {
///     arena: UnsafeCell::new([0x55; ARENA_SIZE]),
///     remaining: AtomicUsize::new(ARENA_SIZE),
/// };
///
/// unsafe impl Sync for SimpleAllocator {}
///
/// unsafe impl GlobalAlloc for SimpleAllocator {
///     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
///         let size = layout.size();
///         let align = layout.align();
///
///         // `Layout` contract forbids making a `Layout` with align=0, or align not power of 2.
///         // So we can safely use a mask to ensure alignment without worrying about UB.
///         let align_mask_to_round_down = !(align - 1);
///
///         if align > MAX_SUPPORTED_ALIGN {
///             return null_mut();
///         }
///
///         let mut allocated = 0;
///         if self
///             .remaining
///             .try_update(Relaxed, Relaxed, |mut remaining| {
///                 if size > remaining {
///                     return None;
///                 }
///                 remaining -= size;
///                 remaining &= align_mask_to_round_down;
///                 allocated = remaining;
///                 Some(remaining)
///             })
///             .is_err()
///         {
///             return null_mut();
///         };
///         unsafe { self.arena.get().cast::<u8>().add(allocated) }
///     }
///     unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
/// }
///
/// fn main() {
///     let _s = format!("allocating a string!");
///     let currently = ALLOCATOR.remaining.load(Relaxed);
///     println!("allocated so far: {}", ARENA_SIZE - currently);
/// }
/// ```
///
/// # The `#[global_allocator]` attribute
///
/// As the example above demonstrates, the `#[global_allocator]` attribute can be used to register a
/// concrete `static` of a type that implements this trait to become *the* global allocator
/// for the current program. That global allocator can be invoked via the functions [`alloc`],
/// [`alloc_zeroed`], [`dealloc`], [`realloc`]). Note, however, that invoking those functions is
/// *not* equivalent to directly invoking the underlying methods on the declared global allocator!
/// Users of the global allocator cannot assume anything about what the allocator does (even if they know which allocator is being used),
/// and implementors of the allocator cannot assume anything about what the program does (even if they know how the allocator is being used).
/// Both can only assume the documented requirements for the respective other party of this contract.
/// This means:
///
/// - Allocation functions may non-deterministically entirely skip the underlying allocator, e.g. if the
///   compiler can show that this allocation can be replaced by a stack variable. The compiler may
///   also merge multiple allocation operations into one, as long as it can also adjust all
///   corresponding deallocation operations accordingly.
/// - An allocation created by invoking [`alloc`], [`alloc_zeroed`], or [`realloc`] has exactly the
///   size and minimum alignment defined by `layout`, even if the underlying allocator makes
///   stronger promises.
/// - An allocation created by invoking [`alloc`], [`alloc_zeroed`], or [`realloc`] can only be
///   freed by invoking [`dealloc`] or [`realloc`]. In particular, passing a pointer to such an
///   allocation directly to the underlying method on [`GlobalAlloc`] is not permitted. Until one of
///   those functions is called, it is undefined behavior to access the memory that backs this
///   allocation with any pointer not derived from the return value of this function (e.g., with
///   internal pointers the allocator might keep around).
/// - The pointer passed to [`dealloc`] or [`realloc`] must have been obtained by invoking [`alloc`],
///   [`alloc_zeroed`], or [`realloc`]. In particular, passing a pointer returned by the underlying
///   methods on [`GlobalAlloc`] is not permitted.
/// - [`alloc`] de-initializes the contents of the allocation before handing it to the user. So even
///   if you control the underlying allocator and know that it explicitly initialized this memory,
///   you cannot rely on it being initialized. For a [`realloc`] that grows an allocation, this
///   applies to the newly allocated part.
/// - [`dealloc`] de-initializes the contents of the allocation before handing it to the allocator.
///   So even if you know that the program previously initialized that memory, the allocator cannot
///   rely on it being initialized. For a [`realloc`] that shrinks an allocation, this applies to
///   the part being removed.
///
/// [`alloc`]: ../../std/alloc/fn.alloc.html
/// [`alloc_zeroed`]: ../../std/alloc/fn.alloc_zeroed.html
/// [`dealloc`]: ../../std/alloc/fn.dealloc.html
/// [`realloc`]: ../../std/alloc/fn.realloc.html
///
/// The first point means that you cannot rely on global allocations actually happening, even if
/// there are explicit global allocations in the source. The optimizer may detect unused global
/// allocations that it can either eliminate entirely or move to the stack and thus never invoke the
/// global 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 global 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.
///
/// # Safety
///
/// The `GlobalAlloc` trait is an `unsafe` trait for a number of reasons, and
/// implementors must ensure that they adhere to these contracts:
///
/// * It is undefined behavior for the allocator to read, write, or deallocate any memory that
///   is *currently allocated*. This memory is owned by the user, the allocator must not touch it.
///
/// * 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.
///
/// * Callers of this trait are allowed to rely on the contracts defined on each method, and
///   implementors must ensure such contracts remain true.
///
/// # 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
#[stable(feature = "global_alloc", since = "1.28.0")]
pub unsafe trait GlobalAlloc {
    /// Allocates memory as described by the given `layout`.
    ///
    /// Returns a pointer to newly-allocated memory,
    /// or null to indicate allocation failure.
    ///
    /// # Safety
    ///
    /// `layout` must have non-zero size. Attempting to allocate for a zero-sized `layout` will
    /// result in undefined behavior.
    ///
    /// (Extension subtraits might provide more specific bounds on
    /// behavior, e.g., guarantee a sentinel address or a null pointer
    /// in response to a zero-size allocation request.)
    ///
    /// The allocated block of memory may or may not be initialized.
    ///
    /// # Errors
    ///
    /// Returning a null pointer indicates that either memory is exhausted
    /// or `layout` does not meet this allocator's size or alignment constraints.
    ///
    /// Implementations are encouraged to return null on memory
    /// exhaustion rather than aborting, but this is not
    /// a strict requirement. (Specifically: it is *legal* to
    /// implement this trait atop an underlying native allocation
    /// library that aborts on memory exhaustion.)
    ///
    /// Clients wishing to abort computation in response to an
    /// allocation error are encouraged to call the [`handle_alloc_error`] function,
    /// rather than directly invoking `panic!` or similar (but note that both may unwind).
    ///
    /// [`handle_alloc_error`]: ../../alloc/alloc/fn.handle_alloc_error.html
    #[stable(feature = "global_alloc", since = "1.28.0")]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8;

    /// Deallocates the block of memory at the given `ptr` pointer with the given `layout`.
    ///
    /// # Safety
    ///
    /// The caller must ensure:
    ///
    /// * `ptr` is a block of memory currently allocated via this allocator and,
    ///
    /// * `layout` is the same layout that was used to allocate that block of
    ///   memory.
    ///
    /// Otherwise the behavior is undefined.
    #[stable(feature = "global_alloc", since = "1.28.0")]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout);

    /// Behaves like `alloc`, but also ensures that the contents
    /// are set to zero before being returned.
    ///
    /// # Safety
    ///
    /// The caller has to ensure that `layout` has non-zero size. Like `alloc`
    /// zero sized `layout` will result in undefined behavior.
    /// However the allocated block of memory is guaranteed to be initialized.
    ///
    /// # Errors
    ///
    /// Returning a null pointer indicates that either memory is exhausted
    /// or `layout` does not meet allocator's size or alignment constraints,
    /// just as in `alloc`.
    ///
    /// Clients wishing to abort computation in response to an
    /// allocation error are encouraged to call the [`handle_alloc_error`] function,
    /// rather than directly invoking `panic!` or similar.
    ///
    /// [`handle_alloc_error`]: ../../alloc/alloc/fn.handle_alloc_error.html
    #[stable(feature = "global_alloc", since = "1.28.0")]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let size = layout.size();
        // SAFETY: the safety contract for `alloc` must be upheld by the caller.
        let ptr = unsafe { self.alloc(layout) };
        if !ptr.is_null() {
            // SAFETY: as allocation succeeded, the region from `ptr`
            // of size `size` is guaranteed to be valid for writes.
            unsafe { ptr::write_bytes(ptr, 0, size) };
        }
        ptr
    }

    /// Shrinks or grows a block of memory to the given `new_size` in bytes.
    /// The block is described by the given `ptr` pointer and `layout`.
    ///
    /// If this returns a non-null pointer, then ownership of the memory block
    /// referenced by `ptr` has been transferred to this allocator.
    /// Any access to the old `ptr` is Undefined Behavior, even if the
    /// allocation remained in-place. The newly returned pointer is the only valid pointer
    /// for accessing this memory now.
    ///
    /// The new memory block is allocated with `layout`,
    /// but with the `size` updated to `new_size` in bytes.
    /// This new layout must be used when deallocating the new memory block with `dealloc`.
    /// The range `0..min(layout.size(), new_size)` of the new memory block is
    /// guaranteed to have the same values as the original block.
    ///
    /// If this method returns null, then ownership of the memory
    /// block has not been transferred to this allocator, and the
    /// contents of the memory block are unaltered.
    ///
    /// # Safety
    ///
    /// The caller must ensure that:
    ///
    /// * `ptr` is allocated via this allocator,
    ///
    /// * `layout` is the same layout that was used
    ///   to allocate that block of memory,
    ///
    /// * `new_size` is greater than zero.
    ///
    /// * `new_size`, when rounded up to the nearest multiple of `layout.align()`,
    ///   does not overflow `isize` (i.e., the rounded value must be less than or
    ///   equal to `isize::MAX`).
    ///
    /// If these are not followed, the behavior is undefined.
    ///
    /// (Extension subtraits might provide more specific bounds on
    /// behavior, e.g., guarantee a sentinel address or a null pointer
    /// in response to a zero-size allocation request.)
    ///
    /// # Errors
    ///
    /// Returns null if the new layout does not meet the size
    /// and alignment constraints of the allocator, or if reallocation
    /// otherwise fails.
    ///
    /// Implementations are encouraged to return null on memory
    /// exhaustion rather than panicking or aborting, but this is not
    /// a strict requirement. (Specifically: it is *legal* to
    /// implement this trait atop an underlying native allocation
    /// library that aborts on memory exhaustion.)
    ///
    /// Clients wishing to abort computation in response to a
    /// reallocation error are encouraged to call the [`handle_alloc_error`] function,
    /// rather than directly invoking `panic!` or similar.
    ///
    /// [`handle_alloc_error`]: ../../alloc/alloc/fn.handle_alloc_error.html
    #[stable(feature = "global_alloc", since = "1.28.0")]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        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: the caller must ensure that `new_layout` is greater than zero.
        let new_ptr = unsafe { self.alloc(new_layout) };
        if !new_ptr.is_null() {
            // SAFETY: the previously allocated block cannot overlap the newly allocated block.
            // The safety contract for `dealloc` must be upheld by the caller.
            unsafe {
                ptr::copy_nonoverlapping(ptr, new_ptr, cmp::min(layout.size(), new_size));
                self.dealloc(ptr, layout);
            }
        }
        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(),
        }
    }
}