fastarena 0.2.0

use core::cmp::Ordering;
use core::hash::{Hash, Hasher};
use core::marker::PhantomData;
use core::mem::{self, ManuallyDrop, MaybeUninit};
use core::ops::{Deref, DerefMut};
use core::ptr::{self, NonNull};

use super::allocator::Arena;

// SAFETY: `ArenaBox<T>` owns a `T` stored in arena memory. Sending the box to
// another thread is equivalent to sending `T`, so we require `T: Send`.
// Likewise for `Sync`. This matches the bounds on `Box<T>`.
unsafe impl<T: Send> Send for ArenaBox<'_, T> {}
unsafe impl<T: Sync> Sync for ArenaBox<'_, T> {}

/// An owned, single-element allocation from an [`Arena`], analogous to
/// [`Box<T>`](std::boxed::Box) but backed by arena memory.
///
/// `ArenaBox` provides single-ownership semantics: when the box is dropped,
/// `T`'s destructor runs immediately. The arena memory itself is reclaimed
/// later by [`Arena::reset`] / [`Arena::rewind`] / arena drop. This makes
/// `ArenaBox` safe to use **without** the `drop-tracking` feature — element
/// destructors are always handled correctly.
///
/// The lifetime `'arena` ties the box to the arena that created it. The
/// arena cannot be reset, rewound, or dropped while any `ArenaBox` from it
/// is alive — this is enforced by the Rust borrow checker.
///
/// # Differences from `Box<T>`
///
/// - **Allocation cost** — O(1) bump-pointer instead of a malloc call.
/// - **Deallocation** — only the destructor runs on `Drop`; the underlying
///   memory is reclaimed in bulk by the arena, not on a per-box basis.
/// - **Lifetimes** — `ArenaBox` cannot outlive its arena.
///
/// # Example
///
/// ```rust
/// use fastarena::{Arena, ArenaBox};
///
/// let mut arena = Arena::new();
/// let mut x: ArenaBox<i32> = arena.alloc_box(42);
/// assert_eq!(*x, 42);
/// *x = 100;
/// assert_eq!(*x, 100);
/// ```
#[repr(transparent)]
pub struct ArenaBox<'arena, T: ?Sized> {
    ptr: NonNull<T>,
    /// Borrows the arena to ensure it outlives the box.
    _marker: PhantomData<&'arena mut T>,
}

impl<'arena, T> ArenaBox<'arena, T> {
    /// Construct a new `ArenaBox<T>` by allocating `val` in the arena.
    ///
    /// Internal — used by [`Arena::alloc_box`].
    #[inline]
    pub(crate) fn new(arena: &'arena mut Arena, val: T) -> Self {
        // SAFETY: We allocate raw memory, write `val` into it, and wrap in
        // ArenaBox. We deliberately do NOT register with the drop registry —
        // ArenaBox::Drop is responsible for running `T`'s destructor.
        if mem::size_of::<T>() == 0 {
            // ZSTs share a dangling but well-aligned address. Drop runs in
            // ArenaBox::Drop just like for non-ZSTs.
            // Forget val to keep ownership semantics — Drop will drop it.
            // Actually, for ZST we must not double-drop: consume val by writing
            // it to the dangling pointer (which is a no-op for ZST) and let
            // ArenaBox::Drop handle it.
            let ptr = NonNull::<T>::dangling();
            // SAFETY: writing a ZST through a dangling-but-aligned pointer is
            // sound — no bytes are touched.
            unsafe { ptr.as_ptr().write(val) };
            return ArenaBox {
                ptr,
                _marker: PhantomData,
            };
        }
        let raw = arena.alloc_raw(mem::size_of::<T>(), mem::align_of::<T>());
        // SAFETY: `raw` points to `size_of::<T>()` aligned bytes inside the arena.
        unsafe {
            let typed = raw.as_ptr().cast::<T>();
            typed.write(val);
            ArenaBox {
                ptr: NonNull::new_unchecked(typed),
                _marker: PhantomData,
            }
        }
    }

    /// Consume the `ArenaBox`, returning the wrapped `T` by value.
    ///
    /// Unlike dropping the box, this does **not** call `T`'s destructor — the
    /// caller becomes responsible for the value.
    ///
    /// The arena memory previously occupied by the value is **not** released
    /// until the arena is reset, rewound, or dropped.
    ///
    /// # Example
    ///
    /// ```rust
    /// use fastarena::Arena;
    ///
    /// let mut arena = Arena::new();
    /// let b = arena.alloc_box(String::from("hello"));
    /// let s: String = b.into_inner();
    /// assert_eq!(s, "hello");
    /// ```
    #[inline]
    #[must_use = "use `drop` to discard the inner value explicitly"]
    pub fn into_inner(self) -> T {
        let this = ManuallyDrop::new(self);
        // SAFETY: `this.ptr` is a valid initialized pointer to `T`. After the
        // read, `this` is forgotten so `Drop` will not run again.
        unsafe { ptr::read(this.ptr.as_ptr()) }
    }

    /// Construct an `ArenaBox` from a raw, fully-initialized pointer to a
    /// value previously allocated in `arena`.
    ///
    /// # Safety
    ///
    /// - `ptr` must point to a fully initialized `T` allocated from `arena`.
    /// - The pointer must not have been registered with the drop registry
    ///   (otherwise a double-drop will occur on reset / rewind).
    /// - No other `ArenaBox`, `&T`, or `&mut T` may alias this allocation.
    #[inline]
    pub unsafe fn from_raw(_arena: &'arena mut Arena, ptr: NonNull<T>) -> Self {
        ArenaBox {
            ptr,
            _marker: PhantomData,
        }
    }

    /// Construct an `ArenaBox` by initializing a raw uninitialized slot.
    ///
    /// This is useful when you want to avoid moving `T` through the stack and
    /// instead construct it in place.
    ///
    /// # Example
    ///
    /// ```rust
    /// use fastarena::{Arena, ArenaBox};
    ///
    /// let mut arena = Arena::new();
    /// let b: ArenaBox<u64> = ArenaBox::new_with(&mut arena, || 0xdead_beef);
    /// assert_eq!(*b, 0xdead_beef);
    /// ```
    #[inline]
    pub fn new_with<F>(arena: &'arena mut Arena, f: F) -> Self
    where
        F: FnOnce() -> T,
    {
        Self::new(arena, f())
    }

    /// Allocate space for a `T` without initializing it.
    ///
    /// The returned box wraps a [`MaybeUninit<T>`]. Use
    /// [`MaybeUninit::write`] to initialize and then call
    /// [`ArenaBox::assume_init`] to convert to `ArenaBox<T>`.
    ///
    /// # Example
    ///
    /// ```rust
    /// use fastarena::{Arena, ArenaBox};
    ///
    /// let mut arena = Arena::new();
    /// let mut slot = ArenaBox::<u64>::new_uninit(&mut arena);
    /// slot.write(42);
    /// // SAFETY: just initialized.
    /// let b = unsafe { ArenaBox::assume_init(slot) };
    /// assert_eq!(*b, 42);
    /// ```
    #[inline]
    pub fn new_uninit(arena: &'arena mut Arena) -> ArenaBox<'arena, MaybeUninit<T>> {
        if mem::size_of::<T>() == 0 {
            return ArenaBox {
                ptr: NonNull::<MaybeUninit<T>>::dangling(),
                _marker: PhantomData,
            };
        }
        let raw = arena.alloc_raw(mem::size_of::<T>(), mem::align_of::<T>());
        ArenaBox {
            ptr: raw.cast::<MaybeUninit<T>>(),
            _marker: PhantomData,
        }
    }
}

impl<'arena, T> ArenaBox<'arena, MaybeUninit<T>> {
    /// Assume the value is fully initialized, converting to `ArenaBox<T>`.
    ///
    /// # Safety
    ///
    /// The caller must ensure the inner `MaybeUninit<T>` has been fully
    /// initialized — failure to do so is undefined behaviour.
    #[inline]
    #[must_use]
    pub unsafe fn assume_init(self) -> ArenaBox<'arena, T> {
        let this = ManuallyDrop::new(self);
        ArenaBox {
            ptr: this.ptr.cast::<T>(),
            _marker: PhantomData,
        }
    }
}

impl<'arena, T: ?Sized> ArenaBox<'arena, T> {
    /// Returns a const raw pointer to the inner value.
    ///
    /// The pointer is valid for the lifetime of the box.
    #[inline]
    pub fn as_ptr(&self) -> *const T {
        self.ptr.as_ptr()
    }

    /// Returns a mutable raw pointer to the inner value.
    #[inline]
    pub fn as_mut_ptr(&mut self) -> *mut T {
        self.ptr.as_ptr()
    }

    /// Consume the box without running `T`'s destructor, returning the raw
    /// pointer for manual reconstruction with [`from_raw`](Self::from_raw).
    #[inline]
    #[must_use = "discarding the raw pointer leaks T's destructor"]
    pub fn into_raw(b: Self) -> NonNull<T> {
        let this = ManuallyDrop::new(b);
        this.ptr
    }

    /// Leak the box, returning a `&mut T` valid for the lifetime of the arena.
    ///
    /// The destructor will not run; the caller can still observe the value
    /// for as long as the arena lives.
    #[inline]
    #[must_use = "leak() returns a reference; discarding it abandons the value"]
    pub fn leak(b: Self) -> &'arena mut T {
        let this = ManuallyDrop::new(b);
        // SAFETY: `this.ptr` is valid for `'arena` (we hold &mut Arena until
        // this box is constructed; no aliases exist).
        unsafe { &mut *this.ptr.as_ptr() }
    }
}

impl<'arena, T: ?Sized> Deref for ArenaBox<'arena, T> {
    type Target = T;

    #[inline]
    fn deref(&self) -> &T {
        // SAFETY: ArenaBox owns the pointee; no other references can exist
        // while we hold &self.
        unsafe { self.ptr.as_ref() }
    }
}

impl<'arena, T: ?Sized> DerefMut for ArenaBox<'arena, T> {
    #[inline]
    fn deref_mut(&mut self) -> &mut T {
        // SAFETY: ArenaBox owns the pointee uniquely; no other references can
        // exist while we hold &mut self.
        unsafe { self.ptr.as_mut() }
    }
}

impl<T: ?Sized> Drop for ArenaBox<'_, T> {
    #[inline]
    fn drop(&mut self) {
        // SAFETY: `self.ptr` points to a valid, initialized `T`. We drop it
        // in place; the arena memory itself remains until the arena is reset
        // or dropped.
        unsafe { ptr::drop_in_place(self.ptr.as_ptr()) };
    }
}

// ---- Trait implementations matching `Box<T>` ----

impl<T: ?Sized + core::fmt::Debug> core::fmt::Debug for ArenaBox<'_, T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        core::fmt::Debug::fmt(&**self, f)
    }
}

impl<T: ?Sized + core::fmt::Display> core::fmt::Display for ArenaBox<'_, T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        core::fmt::Display::fmt(&**self, f)
    }
}

impl<T: ?Sized> core::fmt::Pointer for ArenaBox<'_, T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        core::fmt::Pointer::fmt(&self.ptr.as_ptr(), f)
    }
}

impl<T: ?Sized + PartialEq> PartialEq for ArenaBox<'_, T> {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        PartialEq::eq(&**self, &**other)
    }
}

impl<T: ?Sized + Eq> Eq for ArenaBox<'_, T> {}

impl<T: ?Sized + PartialOrd> PartialOrd for ArenaBox<'_, T> {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        PartialOrd::partial_cmp(&**self, &**other)
    }
}

impl<T: ?Sized + Ord> Ord for ArenaBox<'_, T> {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        Ord::cmp(&**self, &**other)
    }
}

impl<T: ?Sized + Hash> Hash for ArenaBox<'_, T> {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        (**self).hash(state)
    }
}

impl<T: ?Sized> AsRef<T> for ArenaBox<'_, T> {
    #[inline]
    fn as_ref(&self) -> &T {
        self
    }
}

impl<T: ?Sized> AsMut<T> for ArenaBox<'_, T> {
    #[inline]
    fn as_mut(&mut self) -> &mut T {
        self
    }
}

impl<T: ?Sized> core::borrow::Borrow<T> for ArenaBox<'_, T> {
    #[inline]
    fn borrow(&self) -> &T {
        self
    }
}

impl<T: ?Sized> core::borrow::BorrowMut<T> for ArenaBox<'_, T> {
    #[inline]
    fn borrow_mut(&mut self) -> &mut T {
        self
    }
}