stumpalo 0.5.0

#![warn(clippy::pedantic)]
#![doc = include_str!(concat!(env!("OUT_DIR"), "/README-rustdocified.md"))]

#[cfg(any(feature = "allocator-api2", feature = "nightly"))]
mod allocator_api;

#[cfg(nightly)]
use core::iter::TrustedLen;
use core::{
    alloc::Layout,
    cell::Cell,
    ptr::{self, NonNull},
};
use core_alloc::boxed::Box;
use core_alloc::vec::Vec;

type Result<T> = core::result::Result<T, Error>;

use crate::{
    arena_ref::ArenaRef,
    chunk::{ChunkHeader, FIRST_HEADER, HEADER_SIZE, INITIAL_CHUNK_CAPACITY},
    errors::*,
    join_chunk_chains,
};

// INVARIANT: `layout.size()` is a multiple of `layout.align()`, in
// all layouts passed to the internal allocation routines.
//
// Layouts are either constructed from `Layout::new::<T>()` or
// `Layout::array::<T>(len)`, which uphold this invariant, or are
// caller-provided layouts normalized with `Layout::pad_to_align()`.
//
// The Rust language guarantees `size_of::<T>() % align_of::<T>() == 0`
// for all types (<https://doc.rust-lang.org/reference/type-layout.html>).

#[doc(hidden)]
extern crate alloc as core_alloc;

/// A bump allocator whose fast path compiles to as few as six instructions.
///
/// Create one with [`Arena::new()`], then allocate via [`alloc`](Arena::alloc),
/// [`alloc_slice_copy`](Arena::alloc_slice_copy), etc.  Chunks are allocated lazily
/// and freed on drop.  All allocation methods take `&self` (interior mutability
/// via [`Cell`]).
///
// # Layout compatibility with `ArenaRef`
//
// `Arena` and `ArenaRef` have the same `repr(C)` field layout so that
// [`as_arena_ref`](Arena::as_arena_ref) can reinterpret `&Arena` as
// `&ArenaRef<'_>` without undefined behaviour.
#[repr(C)]
pub struct Arena {
    pub(crate) top: Cell<*mut u8>,
    pub(crate) bottom: Cell<*mut ChunkHeader>,
    pub(crate) fresh_chunks: Cell<*mut ChunkHeader>,
    // Padding: matches the size of ArenaRef's `arena_ptr` field so that
    // `as_arena_ref` can cast &Arena -> &ArenaRef without reading past the
    // allocation. The reinterpreted `arena_ptr` is never dereferenced.
    // It would aparently be undefined behaviour to have this, even though
    // the field is never read (according to miri, anyway).
    pub(crate) _padding: *const (),
}

// Arena owns all of its state.
// Calls to the global alloc/dealloc are fine to call from any thread.
unsafe impl Send for Arena {}

impl Default for Arena {
    fn default() -> Self {
        Self::new()
    }
}

impl Arena {
    /// Create an empty arena.  No chunk is allocated until the first `alloc`.
    #[must_use]
    pub fn new() -> Self {
        let header = ptr::from_ref(&FIRST_HEADER).cast_mut();
        let bottom = unsafe { header.add(1) };
        Self {
            top: Cell::new(bottom.cast::<u8>()),
            bottom: Cell::new(bottom),
            fresh_chunks: Cell::new(ptr::null_mut()),
            _padding: ptr::null(),
        }
    }

    /// Create a new arena with the default chunk size.
    /// and allocate a chunk of that size.
    /// See [`with_chunk`](Arena::with_chunk) for a reason why you might want to
    /// do this.
    #[must_use]
    pub fn with_default_chunk() -> Self {
        Self::with_chunk(INITIAL_CHUNK_CAPACITY)
    }

    /// Fallible version of [`with_chunk`](Self::with_chunk).
    #[must_use]
    pub fn try_with_default_chunk() -> Result<Self> {
        Self::try_with_chunk(INITIAL_CHUNK_CAPACITY)
    }

    /// Create an arena, allocating the first chunk upfront.
    /// This can avoid reinforcing the wrong conditional branch
    /// on the first call to `alloc`.
    #[must_use]
    pub fn with_chunk(default_chunk_size: usize) -> Self {
        let header = ChunkHeader::new(default_chunk_size, ptr::null_mut());
        let bottom = unsafe { header.add(1) };
        let top = unsafe { bottom.byte_add((*header).capacity).cast::<u8>() };
        Self {
            top: Cell::new(top),
            bottom: Cell::new(bottom),
            fresh_chunks: Cell::new(ptr::null_mut()),
            _padding: ptr::null(),
        }
    }

    /// Fallible version of [`with_chunk`](Self::with_chunk).
    pub fn try_with_chunk(first_chunk_size: usize) -> Result<Self> {
        let header = ChunkHeader::try_new(first_chunk_size, ptr::null_mut())?;
        let bottom = unsafe { header.add(1) };
        let top = unsafe { bottom.byte_add((*header).capacity).cast::<u8>() };
        Ok(Self {
            top: Cell::new(top),
            bottom: Cell::new(bottom),
            fresh_chunks: Cell::new(ptr::null_mut()),
            _padding: ptr::null(),
        })
    }

    // --- Allocation forwarding methods ---
    // These forward to ArenaRef by reinterpreting &self as &ArenaRef via
    // transmute. Arena and ArenaRef share the same repr(C) layout for their
    // first three fields (top, bottom, fresh_chunks), so treating an Arena
    // as an ArenaRef is sound.  The arena pointer field in the reinterpreted
    // ArenaRef is never read during allocation (only in ArenaRef::Drop,
    // which does not run on references), so dangling bytes there are harmless.

    #[inline(always)]
    pub fn as_arena_ref(&self) -> &ArenaRef<'_> {
        // SAFETY: Arena and ArenaRef have identical #[repr(C)] layouts for
        // the first three fields (top, bottom, fresh_chunks). The arena
        // field is only accessed in ArenaRef::Drop, which does not run on
        // the returned reference.
        unsafe { &*(ptr::from_ref(self).cast::<ArenaRef<'_>>()) }
    }

    #[inline(always)]
    pub fn as_arena_ref_mut(&mut self) -> &mut ArenaRef<'_> {
        // SAFETY: Arena and ArenaRef have identical #[repr(C)] layouts for
        // the first three fields (top, bottom, fresh_chunks). The arena
        // field is only accessed in ArenaRef::Drop, which does not run on
        // the returned reference.
        unsafe { &mut *(ptr::from_mut(self).cast::<ArenaRef<'_>>()) }
    }

    /// Allocate a `T` into the arena. Returns a mutable reference to the copy.
    /// Note that `T`'s destructor will *not* be called on the copy.
    #[inline]
    pub fn alloc<T>(&self, data: T) -> &mut T {
        self.as_arena_ref().alloc(data)
    }

    /// Fallible version of [`alloc`](Self::alloc).
    #[inline]
    pub fn try_alloc<T>(&self, data: T) -> Result<&mut T> {
        self.as_arena_ref().try_alloc(data)
    }

    /// This is a version of [alloc](ArenaRef::alloc), marked as `#[inline(never)]`.
    /// This can be useful if you want to optimize for code size, as it
    /// doesn't emit a new conditional branch, it may even be faster when you have
    /// lots of different allocations in your codepath, as it'll free up branch
    /// predictor cache.
    /// You should default to `alloc`.
    #[inline(never)]
    pub fn alloc_no_inline<T>(&self, data: T) -> &mut T {
        self.alloc(data)
    }

    /// Fallible version of [`alloc_no_inline`](Self::alloc_no_inline).
    #[inline(never)]
    pub fn try_alloc_no_inline<T>(&self, data: T) -> Result<&mut T> {
        self.try_alloc(data)
    }

    /// Allocate a `T` into the arena, as constructed by the closure.
    /// Returns a mutable reference to the allocated value.
    #[inline]
    pub fn alloc_with<F, T>(&self, f: F) -> &mut T
    where
        F: FnOnce() -> T,
    {
        self.as_arena_ref().alloc_with(f)
    }

    /// Fallible version of [`alloc_with`](Self::alloc_with).
    #[inline]
    pub fn try_alloc_with<F, T>(&self, f: F) -> Result<&mut T>
    where
        F: FnOnce() -> T,
    {
        self.as_arena_ref().try_alloc_with(f)
    }

    /// Copy a vector into the arena, consuming it.
    /// This is more efficient than calling `arena.alloc_slice_fill_iter(vec.into_iter())`.
    #[inline]
    pub fn alloc_vec<T>(&self, data: Vec<T>) -> &mut [T] {
        self.as_arena_ref().alloc_vec(data)
    }

    /// Fallible version of [`alloc_vec`](Self::alloc_vec).
    #[inline]
    pub fn try_alloc_vec<T>(&self, data: Vec<T>) -> Result<&mut [T]> {
        self.as_arena_ref().try_alloc_vec(data)
    }

    /// Copy a boxed slice into the arena, consuming it.
    /// This is more efficient than calling `arena.alloc_slice_fill_iter(boxed_slice.into_iter())`.
    #[inline]
    pub fn alloc_boxed_slice<T>(&self, data: Box<[T]>) -> &mut [T] {
        self.as_arena_ref().alloc_boxed_slice(data)
    }

    /// Fallible version of [`alloc_boxed_slice`](Self::alloc_boxed_slice).
    #[inline]
    pub fn try_alloc_boxed_slice<T>(&self, data: Box<[T]>) -> Result<&mut [T]> {
        self.as_arena_ref().try_alloc_boxed_slice(data)
    }

    /// Allocate space for an object with the given [`Layout`].
    ///
    /// The returned pointer points at uninitialized memory.
    ///
    /// # Panics
    ///
    /// Panics if reserving space matching `layout` fails.
    #[inline]
    pub fn alloc_layout(&self, layout: Layout) -> NonNull<[u8]> {
        self.as_arena_ref().alloc_layout(layout)
    }

    /// Fallible version of [`alloc_layout`](Self::alloc_layout).
    #[inline]
    pub fn try_alloc_layout(&self, layout: Layout) -> Result<NonNull<[u8]>> {
        self.as_arena_ref().try_alloc_layout(layout)
    }

    /// Copy a slice into the arena. Returns a mutable reference to the copy.
    #[inline]
    pub fn alloc_slice_copy<T: Copy>(&self, data: &[T]) -> &mut [T] {
        self.as_arena_ref().alloc_slice_copy(data)
    }

    /// Fallible version of [`alloc_slice_copy`](Self::alloc_slice_copy).
    #[inline]
    pub fn try_alloc_slice_copy<T: Copy>(&self, data: &[T]) -> Result<&mut [T]> {
        self.as_arena_ref().try_alloc_slice_copy(data)
    }

    /// Copy a slice literal into the arena. Returns a mutable reference to the copy.
    /// This is faster than [`alloc_slice_copy`](Arena::alloc_slice_copy), as it assumes the size is
    /// available at compile-time.
    /// If you use this on a non-literal slice, (eg. one produced by [leak](Box::leak)),
    /// then this will have slightly worse performance than [`alloc_slice_copy`](Arena::alloc_slice_copy).
    #[inline]
    pub fn alloc_slice_lit_copy<T: Copy>(&self, data: &'static [T]) -> &mut [T] {
        self.as_arena_ref().alloc_slice_lit_copy(data)
    }

    /// Fallible version of [`alloc_slice_lit_copy`](Self::alloc_slice_lit_copy).
    #[inline]
    pub fn try_alloc_slice_lit_copy<T: Copy>(&self, data: &'static [T]) -> Result<&mut [T]> {
        self.as_arena_ref().try_alloc_slice_lit_copy(data)
    }

    /// Copy a slice into the arena via cloning. Returns a mutable reference to the copy.
    #[inline]
    pub fn alloc_slice_clone<T: Clone>(&self, data: &[T]) -> &mut [T] {
        self.as_arena_ref().alloc_slice_clone(data)
    }

    /// Fallible version of [`alloc_slice_clone`](Self::alloc_slice_clone).
    #[inline]
    pub fn try_alloc_slice_clone<T: Clone>(&self, data: &[T]) -> Result<&mut [T]> {
        self.as_arena_ref().try_alloc_slice_clone(data)
    }

    /// Allocate an array of known size, populated by a function.
    #[inline]
    pub fn alloc_sized_slice_with<T, const N: usize, F>(&self, f: F) -> &mut [T; N]
    where
        F: FnOnce() -> [T; N],
    {
        self.as_arena_ref().alloc_sized_slice_with(f)
    }

    /// Fallible version of [`alloc_sized_slice_with`](Self::alloc_sized_slice_with).
    #[inline]
    pub fn try_alloc_sized_slice_with<T, const N: usize, F>(&self, f: F) -> Result<&mut [T; N]>
    where
        F: FnOnce() -> [T; N],
    {
        self.as_arena_ref().try_alloc_sized_slice_with(f)
    }

    /// Allocate an array of known size, populated element-by-element.
    #[inline]
    pub fn alloc_sized_slice_fill_with<T, const N: usize, F>(&self, f: F) -> &mut [T; N]
    where
        F: FnMut(usize) -> T,
    {
        self.as_arena_ref().alloc_sized_slice_fill_with(f)
    }

    /// Fallible version of [`alloc_sized_slice_fill_with`](Self::alloc_sized_slice_fill_with).
    #[inline]
    pub fn try_alloc_sized_slice_fill_with<T, const N: usize, F>(&self, f: F) -> Result<&mut [T; N]>
    where
        F: FnMut(usize) -> T,
    {
        self.as_arena_ref().try_alloc_sized_slice_fill_with(f)
    }

    /// Allocate a slice, populating all slots uniformly by a function.
    ///
    /// # Panics
    ///
    /// Panics if the size of the array would exceed [`isize::MAX`](isize::MAX).
    #[inline]
    pub fn alloc_slice_fill_with<T, F>(&self, len: usize, f: F) -> &mut [T]
    where
        F: FnMut(usize) -> T,
    {
        self.as_arena_ref().alloc_slice_fill_with(len, f)
    }

    /// Fallible version of [`alloc_slice_fill_with`](Self::alloc_slice_fill_with).
    #[inline]
    pub fn try_alloc_slice_fill_with<T, F>(&self, len: usize, f: F) -> Result<&mut [T]>
    where
        F: FnMut(usize) -> T,
    {
        self.as_arena_ref().try_alloc_slice_fill_with(len, f)
    }

    /// Allocate a slice, populating all slots uniformly by a fallible function.
    ///
    /// # Errors
    ///
    /// Returns the first error returned by `f`.
    /// In case of error, the arena will revert to its state before this call.
    ///
    /// # Panics
    ///
    /// Panics if the size of the array would exceed [`isize::MAX`](isize::MAX).
    #[inline]
    pub fn alloc_slice_try_fill_with<T, E, F>(
        &self,
        len: usize,
        f: F,
    ) -> core::result::Result<&mut [T], E>
    where
        F: FnMut(usize) -> core::result::Result<T, E>,
    {
        self.as_arena_ref().alloc_slice_try_fill_with(len, f)
    }

    /// Copy a fixed-size array into the arena. Returns a mutable reference.
    #[inline]
    pub fn alloc_sized_slice_copy<T: Copy, const N: usize>(&self, data: &[T; N]) -> &mut [T; N] {
        self.as_arena_ref().alloc_sized_slice_copy(data)
    }

    /// Fallible version of [`alloc_sized_slice_copy`](Self::alloc_sized_slice_copy).
    #[inline]
    pub fn try_alloc_sized_slice_copy<T: Copy, const N: usize>(
        &self,
        data: &[T; N],
    ) -> Result<&mut [T; N]> {
        self.as_arena_ref().try_alloc_sized_slice_copy(data)
    }

    /// Copy a fixed-size array into the arena via cloning.
    #[inline]
    pub fn alloc_sized_slice_clone<T: Clone, const N: usize>(&self, data: &[T; N]) -> &mut [T; N] {
        self.as_arena_ref().alloc_sized_slice_clone(data)
    }

    /// Fallible version of [`alloc_sized_slice_clone`](Self::alloc_sized_slice_clone).
    #[inline]
    pub fn try_alloc_sized_slice_clone<T: Clone, const N: usize>(
        &self,
        data: &[T; N],
    ) -> Result<&mut [T; N]> {
        self.as_arena_ref().try_alloc_sized_slice_clone(data)
    }

    /// Allocate a fixed-size array filled with [Default] values.
    #[inline]
    pub fn alloc_sized_slice_fill_default<T: Default, const N: usize>(&self) -> &mut [T; N] {
        self.as_arena_ref().alloc_sized_slice_fill_default()
    }

    /// Fallible version of [`alloc_sized_slice_fill_default`](Self::alloc_sized_slice_fill_default).
    #[inline]
    pub fn try_alloc_sized_slice_fill_default<T: Default, const N: usize>(
        &self,
    ) -> Result<&mut [T; N]> {
        self.as_arena_ref().try_alloc_sized_slice_fill_default()
    }

    /// Allocate a slice filled with [Default] values.
    ///
    /// # Panics
    ///
    /// Panics if the size of the array would exceed [`isize::MAX`](isize::MAX).
    #[inline]
    pub fn alloc_slice_fill_default<T: Default>(&self, len: usize) -> &mut [T] {
        self.as_arena_ref().alloc_slice_fill_default(len)
    }

    /// Fallible version of [`alloc_slice_fill_default`](Self::alloc_slice_fill_default).
    #[inline]
    pub fn try_alloc_slice_fill_default<T: Default>(&self, len: usize) -> Result<&mut [T]> {
        self.as_arena_ref().try_alloc_slice_fill_default(len)
    }

    /// Allocate a slice from an iterator with a known size.
    /// Returns a mutable reference to the allocated slice.
    ///
    /// # Panics
    ///
    /// Panics if the iterator doesn't have as many elements as its `.len()` promises.
    /// Panics if the size of the array would exceed [`isize::MAX`](isize::MAX).
    #[inline]
    pub fn alloc_slice_fill_iter<T, I>(&self, iter: I) -> &mut [T]
    where
        I: IntoIterator<Item = T>,
        <I as IntoIterator>::IntoIter: ExactSizeIterator,
    {
        self.as_arena_ref().alloc_slice_fill_iter(iter)
    }

    /// Fallible version of [`alloc_slice_fill_iter`](Self::alloc_slice_fill_iter).
    #[inline]
    pub fn try_alloc_slice_fill_iter<T, I>(&self, iter: I) -> Result<&mut [T]>
    where
        I: IntoIterator<Item = T>,
        <I as IntoIterator>::IntoIter: ExactSizeIterator,
    {
        self.as_arena_ref().try_alloc_slice_fill_iter(iter)
    }

    /// Allocate a slice from an iterator with a trusted known size.
    /// Returns a mutable reference to the allocated slice.
    ///
    /// # Panics
    ///
    /// Panics if the iterator doesn't have as many elements as its `.len()` promises.
    /// Panics if the size of the array would exceed [`isize::MAX`](isize::MAX).
    #[inline]
    #[cfg(nightly)]
    pub fn alloc_slice_fill_iter_trusted<T, I>(&self, iter: I) -> &mut [T]
    where
        I: IntoIterator<Item = T>,
        <I as IntoIterator>::IntoIter: TrustedLen,
    {
        self.as_arena_ref().alloc_slice_fill_iter_trusted(iter)
    }

    /// Copy a string into the arena. Returns a mutable reference to the copy.
    #[inline]
    pub fn alloc_str(&self, s: &str) -> &mut str {
        self.as_arena_ref().alloc_str(s)
    }

    /// Fallible version of [`alloc_str`](Self::alloc_str).
    #[inline]
    pub fn try_alloc_str(&self, s: &str) -> Result<&mut str> {
        self.as_arena_ref().try_alloc_str(s)
    }

    /// Copy a string literal into the arena. Returns a mutable reference to the copy.
    /// This is faster than [`alloc_str`](Arena::alloc_str), as it assumes the size is
    /// available at compile-time.
    /// If you use this on a non-literal string, (eg. one produced by [leak](String::leak)),
    /// then this will have slightly worse performance than [`alloc_str`](Arena::alloc_str).
    #[inline]
    pub fn alloc_str_lit(&self, s: &'static str) -> &mut str {
        self.as_arena_ref().alloc_str_lit(s)
    }

    /// Fallible version of [`alloc_str_lit`](Self::alloc_str_lit).
    #[inline]
    pub fn try_alloc_str_lit(&self, s: &'static str) -> Result<&mut str> {
        self.as_arena_ref().try_alloc_str_lit(s)
    }

    /// Counts the size of all chunks in this arena, whether those chunks are
    /// in use or not. Doesn't include chunk metadata.
    pub fn allocated_bytes(&self) -> usize {
        self.as_arena_ref().allocated_bytes()
    }

    /// Counts the size of all chunks in this arena, whether those chunks are
    /// in use or not. Includes chunk metadata.
    pub fn allocated_bytes_including_metadata(&self) -> usize {
        self.as_arena_ref().allocated_bytes_including_metadata()
    }

    /// Returns the capacity of the current chunk.
    pub fn chunk_capacity(&self) -> usize {
        self.as_arena_ref().chunk_capacity()
    }

    /// Returns whether the arena has any fresh chunks allocated.
    pub fn has_fresh_chunks(&self) -> bool {
        self.as_arena_ref().has_fresh_chunks()
    }

    fn header(&self) -> *mut ChunkHeader {
        let ptr = self.bottom.get();
        unsafe { ptr.sub(1) }
    }

    /// Clear the arena, invalidating all references to arena-allocated data
    /// at compile time. Chunks that have already been allocated are reused
    /// in subsequent allocations.
    /// If you want to actually free those chunks, just recreate the arena.
    pub fn clear(&mut self) {
        let current = self.header();
        let capacity = unsafe { (*current).capacity };
        // Guard against the FIRST_HEADER sentinel (capacity 0), which
        // lives in read-only static memory and must not be mutated.
        if capacity == 0 {
            return;
        }

        // Reset top to end of current chunk
        self.top
            .set(unsafe { (current.add(1).cast::<u8>()).byte_add(capacity) });
        // Move used chain to fresh_chunks
        let used_chain = unsafe { (*current).next };
        unsafe { (*current).next = ptr::null_mut() };
        self.fresh_chunks
            .set(join_chunk_chains(self.fresh_chunks.get(), used_chain));
    }

    /// Free all chunks in the 'fresh chunk' list.
    /// Fresh chunks only occur if you're using scoped arenas (see [`with_scope`](Self::with_scope)),
    /// or if you've called [clear](Self::clear).
    pub fn free_fresh_chunks(&mut self) {
        unsafe { ChunkHeader::free_chunk_chain(self.fresh_chunks.get()) };
    }

    /// Alias for [clear](Arena::clear).
    pub fn reset(&mut self) {
        self.clear();
    }

    /// Run a closure with a sub-scoped arena.  References returned by the
    /// closure are bounded by the closure's scope (not the parent arena's
    /// lifetime), and the arena state is reverted on return (chunks are
    /// reused, not freed).
    ///
    /// Takes `&mut self`, making references from `alloc` are invalidated.
    /// This is intentional. If you need to use outer allocations after
    /// a scope closes, use an [`ArenaRef`], by calling
    /// [`as_arena_ref_mut`](Arena::as_arena_ref_mut).
    pub fn with_scope<'brand, F, R>(&'brand mut self, f: F) -> R
    where
        F: for<'scope> FnOnce(&'scope mut ArenaRef<'scope>) -> R,
    {
        unsafe { self.as_arena_ref_mut().with_scope_unchecked(f) }
    }
}

impl Drop for Arena {
    fn drop(&mut self) {
        unsafe {
            let bottom = self.bottom.get();
            if !bottom.is_null() {
                let current = bottom.byte_sub(HEADER_SIZE).cast::<ChunkHeader>();
                ChunkHeader::free_chunk_chain(current);
            }
            ChunkHeader::free_chunk_chain(self.fresh_chunks.get());
        }
    }
}