multitude 0.3.2

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

//! Safe `Arena`-level reservation primitives.
//!
//! These helpers wrap [`ChunkMutator::try_alloc_uninit*`] with the
//! lifetime promotion that allows the resulting ticket to be consumed
//! into a `&'arena mut T` reference at the public API boundary.
//!
//! # Why the promotion is sound
//!
//! Local chunks: when [`Arena::refill_local`] rotates a chunk out, it
//! pushes the retired [`ChunkMutator`] into `Arena::retired_local`,
//! where it holds its `+1` refcount. The vector is cleared only on
//! `&mut self` paths (`reset`, `Drop`). Therefore every slot ever
//! reserved in a local chunk remains live for the entire `&Arena`
//! borrow lifetime — even if the chunk has been rotated out.
//!
//! Shared chunks: by design these only produce `Arc<T>` smart pointers,
//! not raw `&mut T`. Each `Arc` independently keeps its hosting chunk
//! alive via the chunk's atomic refcount, so the `&Arena` lifetime
//! promotion isn't relied upon for shared reservations either. The
//! `try_reserve_shared*` helpers still rebind the ticket lifetime for
//! API symmetry; callers wrap the resulting reference in an `Arc`
//! before exposing it.
//!
//! All `unsafe` related to lifetime extension lives in this single
//! module; the public `alloc_*.rs` files contain no `unsafe` blocks.

use core::ptr::NonNull;

use allocator_api2::alloc::Allocator;

use super::Arena;
use crate::internal::shared_chunk::SharedChunk;
use crate::internal::uninit::{Uninit, UninitDrop};

impl<A: Allocator + Clone> Arena<A> {
    /// Try to reserve uninitialized storage for one `T` in the current
    /// local chunk. Returns a [`Uninit`] ticket whose lifetime is bound
    /// to `&self`.
    #[inline(always)]
    #[cfg_attr(test, mutants::skip)] // body→None ⇒ refill spin (OOM)
    pub(crate) fn try_reserve_local<T>(&self) -> Option<Uninit<'_, T>> {
        let ticket = self.current_local().try_alloc_uninit::<T>()?;
        // SAFETY: the chunk that hosts this slot is retained for the
        // full `&Arena` borrow lifetime; see module-level rationale.
        Some(unsafe { ticket.rebind() })
    }

    /// Try to reserve uninitialized storage for one `T` plus a drop
    /// entry slot in the current local chunk.
    #[inline(always)]
    #[cfg_attr(test, mutants::skip)] // see `try_reserve_local`
    pub(crate) fn try_reserve_local_with_drop<T>(&self) -> Option<UninitDrop<'_, T>> {
        let ticket = self.current_local().try_alloc_uninit_with_drop::<T>()?;
        // SAFETY: see module-level rationale.
        Some(unsafe { ticket.rebind() })
    }

    /// Try to reserve uninitialized storage for `len` consecutive `T`s
    /// in the current local chunk.
    #[inline(always)]
    #[cfg_attr(test, mutants::skip)] // see `try_reserve_local`
    pub(crate) fn try_reserve_local_slice<T>(&self, len: usize) -> Option<Uninit<'_, [T]>> {
        let ticket = self.current_local().try_alloc_uninit_slice::<T>(len)?;
        // SAFETY: see module-level rationale.
        Some(unsafe { ticket.rebind() })
    }

    /// Like [`Self::try_reserve_local_slice`] but takes the precomputed
    /// byte size; the slice-copy/clone fast paths hold an existing
    /// `&[T]` and compute `size_of_val(src)` once outside the refill
    /// loop, sparing the inner reservation a `checked_mul` overflow
    /// guard.
    ///
    /// # Safety
    ///
    /// `size` must equal `size_of::<T>() * len` (without overflow).
    #[inline(always)]
    #[cfg_attr(test, mutants::skip)] // see `try_reserve_local`
    pub(crate) fn try_reserve_local_slice_with_size<T>(&self, len: usize, size: usize) -> Option<Uninit<'_, [T]>> {
        // SAFETY: forwarded to the caller.
        let ticket = unsafe { self.current_local().try_alloc_uninit_slice_with_size::<T>(len, size) }?;
        // SAFETY: see module-level rationale.
        Some(unsafe { ticket.rebind() })
    }

    /// Try to reserve `len` consecutive bytes in the current local chunk.
    /// Byte-slice fast path that skips alignment math and overflow checks.
    #[inline(always)]
    #[cfg_attr(test, mutants::skip)] // see `try_reserve_local`
    pub(crate) fn try_reserve_local_bytes(&self, len: usize) -> Option<Uninit<'_, [u8]>> {
        let ticket = self.current_local().try_alloc_bytes(len)?;
        // SAFETY: see module-level rationale.
        Some(unsafe { ticket.rebind() })
    }

    /// Try to reserve uninitialized storage for `len` consecutive `T`s
    /// plus a drop entry slot in the current local chunk.
    #[inline(always)]
    #[cfg_attr(test, mutants::skip)] // see `try_reserve_local`
    pub(crate) fn try_reserve_local_slice_with_drop<T>(&self, len: usize) -> Option<UninitDrop<'_, [T]>> {
        let ticket = self.current_local().try_alloc_uninit_slice_with_drop::<T>(len)?;
        // SAFETY: see module-level rationale.
        Some(unsafe { ticket.rebind() })
    }

    /// Try to reserve uninitialized storage for one `T` in the current
    /// shared chunk. The returned ticket is consumed by the caller
    /// before the chunk's refcount is incremented for the smart pointer.
    /// The chunk pointer is returned alongside the ticket so callers
    /// can take a +1 refcount on it without re-asserting that the
    /// current mutator owns a chunk.
    #[inline(always)]
    #[cfg_attr(test, mutants::skip)] // body→None ⇒ refill spin (OOM); same for the helpers below
    pub(crate) fn try_reserve_shared<T>(&self) -> Option<(Uninit<'_, T>, NonNull<SharedChunk<A>>)> {
        let mutator = self.current_shared();
        let ticket = mutator.try_alloc_uninit::<T>()?;
        // SAFETY: `try_alloc_uninit` returning `Some` proves the
        // mutator owns a chunk; `rebind` is sound per the module-level
        // rationale (the chunk is retained across the `&Arena` borrow).
        Some(unsafe { (ticket.rebind(), mutator.chunk_ptr_unchecked()) })
    }

    /// Try to reserve uninitialized storage for `len` consecutive `T`s
    /// in the current shared chunk, taking the precomputed payload byte
    /// size; the slice-copy fast paths hold an existing `&[T]` and
    /// compute `size_of_val(src)` once outside the refill loop, sparing
    /// the inner reservation a `checked_mul` overflow guard.
    ///
    /// Includes a thin-pointer DST length prefix immediately before
    /// the payload — see [`ChunkMutator::try_alloc_uninit_slice_prefixed`].
    ///
    /// # Safety
    ///
    /// `payload_bytes` must equal `size_of::<T>() * len` (without
    /// overflow).
    #[inline(always)]
    #[cfg_attr(test, mutants::skip)] // see `try_reserve_shared`
    #[allow(
        clippy::type_complexity,
        reason = "ticket + chunk-ptr tuple is the natural shape; type alias would obscure rather than clarify"
    )]
    pub(crate) unsafe fn try_reserve_shared_slice_with_size<T>(
        &self,
        len: usize,
        payload_bytes: usize,
    ) -> Option<(Uninit<'_, [T]>, NonNull<SharedChunk<A>>)> {
        let mutator = self.current_shared();
        // SAFETY: forwarded to the caller.
        let ticket = unsafe { mutator.try_alloc_uninit_slice_prefixed_with_size::<T>(len, payload_bytes) }?;
        // SAFETY: see `try_reserve_shared`.
        Some(unsafe { (ticket.rebind(), mutator.chunk_ptr_unchecked()) })
    }

    /// Try to reserve storage for one strong-prefixed `Arc<T>` value in
    /// the current shared chunk. The returned ticket addresses the
    /// payload (the strong count is already initialized to `1`).
    #[inline(always)]
    #[cfg_attr(test, mutants::skip)] // see `try_reserve_shared`
    pub(crate) fn try_reserve_arc_value<T>(&self) -> Option<(Uninit<'_, T>, NonNull<SharedChunk<A>>)> {
        let (ticket, chunk) = self.current_shared().try_alloc_arc_value::<T>()?;
        // SAFETY: see `try_reserve_shared`.
        Some(unsafe { (ticket.rebind(), chunk) })
    }

    /// Slice form of [`Self::try_reserve_arc_value`]: reserves a strong
    /// prefix, slice-length metadata, and `len` `T`s.
    #[inline(always)]
    #[cfg_attr(test, mutants::skip)] // see `try_reserve_shared`
    #[allow(
        clippy::type_complexity,
        reason = "ticket + chunk-ptr tuple is the natural shape; type alias would obscure rather than clarify"
    )]
    pub(crate) fn try_reserve_arc_slice<T>(&self, len: usize) -> Option<(Uninit<'_, [T]>, NonNull<SharedChunk<A>>)> {
        let (ticket, chunk) = self.current_shared().try_alloc_arc_slice::<T>(len)?;
        // SAFETY: see `try_reserve_shared`.
        Some(unsafe { (ticket.rebind(), chunk) })
    }

    /// Like [`Self::try_reserve_arc_slice`] but takes the precomputed
    /// payload byte size (held by callers with a live `&[T]`).
    ///
    /// # Safety
    ///
    /// `payload_bytes` must equal `size_of::<T>() * len` (without overflow).
    #[inline(always)]
    #[cfg_attr(test, mutants::skip)] // see `try_reserve_shared`
    #[allow(
        clippy::type_complexity,
        reason = "ticket + chunk-ptr tuple is the natural shape; type alias would obscure rather than clarify"
    )]
    pub(crate) unsafe fn try_reserve_arc_slice_with_size<T>(
        &self,
        len: usize,
        payload_bytes: usize,
    ) -> Option<(Uninit<'_, [T]>, NonNull<SharedChunk<A>>)> {
        // SAFETY: forwarded to the caller.
        let (ticket, chunk) = unsafe { self.current_shared().try_alloc_arc_slice_with_size::<T>(len, payload_bytes) }?;
        // SAFETY: see `try_reserve_shared`.
        Some(unsafe { (ticket.rebind(), chunk) })
    }
}