use crate::errors::*;
#[cfg(nightly)]
use core::iter::TrustedLen;
use core::{
    alloc::Layout,
    array,
    cell::Cell,
    hint::cold_path,
    marker::PhantomData,
    mem,
    ptr::{self, NonNull},
};
use core_alloc::alloc::dealloc;
use core_alloc::boxed::Box;
use core_alloc::vec::Vec;

use crate::{
    Arena,
    chunk::{CHUNK_ALIGN, ChunkHeader, HEADER_SIZE, INITIAL_CHUNK_CAPACITY},
    count_bytes_allocated_in_chunk_seq, count_bytes_allocated_in_chunk_seq_with_meta,
    limits::{MIN_TOP, SAFE_SIZE},
};

/// An arena allocator reference, tied to the lifetime of its parent
/// [`Arena`] or [`ArenaRef`].
/// The only time you should need to use this is when using [`with_scope`](Arena::with_scope).
/// It provides mostly the same API as `Arena`.
#[repr(C)]
pub struct ArenaRef<'a> {
    // pointer to the first unfree byte
    top: Cell<*mut u8>,
    // pointer to the current chunk (after the header)
    bottom: Cell<*mut ChunkHeader>,
    // pointer to the next reusable chunk
    fresh_chunks: Cell<*mut ChunkHeader>,
    // Raw pointer to the parent Arena.  Only dereferenced in Drop (and only
    // by real ArenaRef values, not by the reinterpreted view from
    // Arena::as_arena_ref).
    arena_ptr: *const Arena,
    // Lifetime marker for the borrow-checker.
    _marker: PhantomData<&'a ()>,
}

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

// --- Public API ---

// Note that there is no clear() method, as that would not invalidate references allocated through
// the arena...
impl<'a> ArenaRef<'a> {
    // For T: u32, on x86_64, the entire fast path compiles down to:
    // ```asm
    // alloc_u32:
    //   mov     rcx, qword ptr [rdi]      # rcx = self.top (current bump pointer)
    //   and     rcx, -4                   # Align top down to 4-byte boundary (top & !(align-1))
    //   lea     rax, [rcx - 4]            # rax = aligned_top - sizeof(u32); allocated address
    //   cmp     rax, qword ptr [rdi + 8]  # if new_top < self.bottom.get() (out of room):
    //   jb      ArenaRef::alloc_more      #    take the slow path
    //   mov     qword ptr [rdi], rax      # self.top = new_top
    //   mov     dword ptr [rcx - 4], esi  # write the u32 to allocated slot
    //   ret                               # return the &mut u32 in rax
    // ```
    //
    // To get this output on compiler explorer, mark this function
    // as #[inline(never)], and compile with `-O -C relocation-model=static`.
    //
    // When there are multiple allocations, the fast path can go as low as six instructions
    // per allocation.
    //
    /// Copy 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) -> &'a mut T {
        self.alloc_with(|| data)
    }

    /// Fallible version of [`alloc`](Self::alloc).
    pub fn try_alloc<T>(&self, data: T) -> Result<&'a mut T> {
        self.try_alloc_with(|| 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) -> &'a 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<&'a mut T> {
        self.try_alloc(data)
    }

    /// Allocate a `T` into the arena, constructed by the closure.
    /// Returns a mutable reference to the allocated value.
    #[inline]
    pub fn alloc_with<F, T>(&self, f: F) -> &'a mut T
    where
        F: FnOnce() -> T,
    {
        let layout = Layout::new::<T>();
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, true);

        // We can potentially create a jbe instruction directly to the
        // slow path, but to do so when this is inlined, we need to
        // put the parameter in a register, or on the stack.
        // This is slow for larger sizes, so we cap this behaviour at 16 bytes.
        let mut ptr: NonNull<T> = if layout.size() <= 16 {
            if take_slow_path {
                cold_path();
                self.alloc_slow_with(f)
            } else {
                unsafe {
                    let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                    debug_assert_eq!(new_top as usize % layout.align(), 0);
                    debug_assert!(new_top as usize + layout.size() <= top as usize);
                    let ptr: *mut T = new_top.cast::<T>();
                    ptr::write(ptr, f());
                    self.top.set(new_top);
                    NonNull::new_unchecked(ptr)
                }
            }
        } else {
            unsafe {
                let mut ptr: NonNull<T> = if take_slow_path {
                    cold_path();
                    self.alloc_slow_raw::<T>()
                } else {
                    let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                    debug_assert_eq!(new_top as usize % layout.align(), 0);
                    debug_assert!(new_top as usize + layout.size() <= top as usize);
                    self.top.set(new_top);
                    NonNull::new_unchecked(new_top.cast::<T>())
                };
                ptr::write(ptr.as_mut(), f());
                ptr
            }
        };
        unsafe { ptr.as_mut() }
    }

    /// Fallible version of [`alloc_with`](Self::alloc_with).
    #[inline]
    pub fn try_alloc_with<F, T>(&self, f: F) -> Result<&'a mut T>
    where
        F: FnOnce() -> T,
    {
        let layout = Layout::new::<T>();
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, true);

        // We can potentially create a jbe instruction directly to the
        // slow path, but to do so when this is inlined, we need to
        // put the parameter in a register, or on the stack.
        // This is slow for larger sizes, so we cap this behaviour at 16 bytes.
        let mut ptr: NonNull<T> = if layout.size() <= 16 {
            if take_slow_path {
                cold_path();
                self.try_alloc_slow_with(f)?
            } else {
                unsafe {
                    let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                    debug_assert_eq!(new_top as usize % layout.align(), 0);
                    debug_assert!(new_top as usize + layout.size() <= top as usize);
                    let ptr: *mut T = new_top.cast::<T>();
                    ptr::write(ptr, f());
                    self.top.set(new_top);
                    NonNull::new_unchecked(ptr)
                }
            }
        } else {
            unsafe {
                let mut ptr: NonNull<T> = if take_slow_path {
                    cold_path();
                    self.try_alloc_slow_raw::<T>()?
                } else {
                    let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                    debug_assert_eq!(new_top as usize % layout.align(), 0);
                    debug_assert!(new_top as usize + layout.size() <= top as usize);
                    self.top.set(new_top);
                    NonNull::new_unchecked(new_top.cast::<T>())
                };
                ptr::write(ptr.as_mut(), f());
                ptr
            }
        };
        Ok(unsafe { ptr.as_mut() })
    }

    /// 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]> {
        let len = layout.size();
        let layout = layout.pad_to_align();
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        // Pessimistically use checked arithmetic here, because we can't guarantee there's
        // any size/alignment information available at compile-time.
        // Using the standard need_slow_path function would result in extra checks, without
        // the compile-time info...
        // If you do know this information at compile-time. I'd love to hear your
        // use-case for this function, and would be happy to add an alternative.
        let take_slow_path = Self::need_slow_path_checked(layout, top, bottom);

        let ptr = unsafe {
            if take_slow_path {
                cold_path();
                self.alloc_layout_slow(layout)
            } else {
                let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                self.top.set(new_top);
                debug_assert_eq!(new_top as usize % layout.align(), 0);
                debug_assert!(new_top as usize + layout.size() <= top as usize);
                NonNull::new_unchecked(new_top)
            }
        };

        NonNull::slice_from_raw_parts(ptr, len)
    }

    /// Fallible version of [`alloc_with`](Self::alloc_with).
    #[inline]
    pub fn try_alloc_layout(&self, layout: Layout) -> Result<NonNull<[u8]>> {
        let len = layout.size();
        let layout = layout.pad_to_align();
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        // Pessimistically use checked arithmetic here, because we can't guarantee there's
        // any size/alignment information available at compile-time.
        // Using the standard need_slow_path function would result in extra checks, without
        // the compile-time info...
        // If you do know this information at compile-time. I'd love to hear your
        // use-case for this function, and would be happy to add an alternative.
        let take_slow_path = Self::need_slow_path_checked(layout, top, bottom);

        let ptr = unsafe {
            if take_slow_path {
                cold_path();
                self.try_alloc_layout_slow(layout)?
            } else {
                let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                self.top.set(new_top);
                debug_assert_eq!(new_top as usize % layout.align(), 0);
                debug_assert!(new_top as usize + layout.size() <= top as usize);
                NonNull::new_unchecked(new_top)
            }
        };

        Ok(NonNull::slice_from_raw_parts(ptr, len))
    }

    /// Copy a slice into the arena. Returns a mutable reference to the copy.
    #[inline]
    pub fn alloc_slice_copy<T: Copy>(&self, data: &[T]) -> &'a mut [T] {
        let len = data.len();
        let layout = Layout::for_value(data);
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, false);

        unsafe {
            let ptr: *mut T = if take_slow_path {
                cold_path();
                self.alloc_layout_slow(layout).cast().as_ptr()
            } else {
                let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                debug_assert_eq!(new_top as usize % layout.align(), 0);
                debug_assert!(new_top as usize + layout.size() <= top as usize);
                self.top.set(new_top);
                new_top.cast::<T>()
            };
            ptr::copy_nonoverlapping(data.as_ptr(), ptr, len);
            core::slice::from_raw_parts_mut(ptr, len)
        }
    }

    /// Fallible version of [`alloc_slice_copy`](Self::alloc_slice_copy).
    #[inline]
    pub fn try_alloc_slice_copy<T: Copy>(&self, data: &[T]) -> Result<&'a mut [T]> {
        let len = data.len();
        let layout = Layout::for_value(data);
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, false);

        unsafe {
            let ptr: *mut T = if take_slow_path {
                cold_path();
                self.try_alloc_layout_slow(layout)?.cast().as_ptr()
            } else {
                let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                debug_assert_eq!(new_top as usize % layout.align(), 0);
                debug_assert!(new_top as usize + layout.size() <= top as usize);
                self.top.set(new_top);
                new_top.cast::<T>()
            };
            ptr::copy_nonoverlapping(data.as_ptr(), ptr, len);
            Ok(core::slice::from_raw_parts_mut(ptr, len))
        }
    }

    /// Copy a slice literal into the arena. Returns a mutable reference to the copy.
    /// This is faster than [`alloc_slice_copy`](ArenaRef::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`](ArenaRef::alloc_slice_copy).
    #[inline]
    pub fn alloc_slice_lit_copy<T: Copy>(&self, data: &'static [T]) -> &'a mut [T] {
        let len = data.len();
        let layout = Layout::for_value(data);
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, true);

        unsafe {
            let ptr: *mut T = if take_slow_path {
                cold_path();
                self.alloc_layout_slow(layout).cast().as_ptr()
            } else {
                let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                debug_assert_eq!(new_top as usize % layout.align(), 0);
                debug_assert!(new_top as usize + layout.size() <= top as usize);
                self.top.set(new_top);
                new_top.cast::<T>()
            };
            ptr::copy_nonoverlapping(data.as_ptr(), ptr, len);
            core::slice::from_raw_parts_mut(ptr, len)
        }
    }

    /// Fallible version of [`alloc_slice_copy`](Self::alloc_slice_copy).
    #[inline]
    pub fn try_alloc_slice_lit_copy<T: Copy>(&self, data: &'static [T]) -> Result<&'a mut [T]> {
        let len = data.len();
        let layout = Layout::for_value(data);
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, true);

        unsafe {
            let ptr: *mut T = if take_slow_path {
                cold_path();
                self.try_alloc_layout_slow(layout)?.cast().as_ptr()
            } else {
                let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                debug_assert_eq!(new_top as usize % layout.align(), 0);
                debug_assert!(new_top as usize + layout.size() <= top as usize);
                self.top.set(new_top);
                new_top.cast::<T>()
            };
            ptr::copy_nonoverlapping(data.as_ptr(), ptr, len);
            Ok(core::slice::from_raw_parts_mut(ptr, len))
        }
    }

    /// Clone a slice into the arena. Returns a mutable reference to the copy.
    #[inline]
    pub fn alloc_slice_clone<T: Clone>(&self, data: &[T]) -> &'a mut [T] {
        self.alloc_slice_fill_with(data.len(), |n| data[n].clone())
    }

    /// Fallible version of [`alloc_slice_clone`](Self::alloc_slice_clone).
    #[inline]
    pub fn try_alloc_slice_clone<T: Clone>(&self, data: &[T]) -> Result<&'a mut [T]> {
        self.try_alloc_slice_fill_with(data.len(), |n| data[n].clone())
    }

    /// 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) -> &'a mut [T; N]
    where
        F: FnOnce() -> [T; N],
    {
        let layout = Layout::new::<[T; N]>();
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, true);

        // We can potentially create a jbe instruction directly to the
        // slow path, but to do so when this is inlined, we need to
        // put the parameter in a register, or on the stack.
        // This is slow for larger sizes, so we cap this behaviour at 16 bytes.
        let mut ptr: NonNull<[T; N]> = if layout.size() <= 16 {
            unsafe {
                if take_slow_path {
                    cold_path();
                    self.alloc_slow_with::<_, [T; N]>(f)
                } else {
                    let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                    debug_assert_eq!(new_top as usize % layout.align(), 0);
                    debug_assert!(new_top as usize + layout.size() <= top as usize);
                    let ptr = (new_top.cast::<[T; N]>()).as_mut_unchecked();
                    let data = mem::ManuallyDrop::new(f());
                    ptr::copy_nonoverlapping(data.as_ptr(), ptr.as_mut_ptr(), data.len());
                    self.top.set(new_top);
                    NonNull::new_unchecked(new_top.cast::<[T; N]>())
                }
            }
        } else {
            unsafe {
                let ptr: NonNull<[T; N]> = if take_slow_path {
                    cold_path();
                    self.alloc_slow_raw::<[T; N]>()
                } else {
                    let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                    debug_assert_eq!(new_top as usize % layout.align(), 0);
                    debug_assert!(new_top as usize + layout.size() <= top as usize);
                    self.top.set(new_top);
                    NonNull::new_unchecked(new_top.cast::<[T; N]>())
                };
                let data = mem::ManuallyDrop::new(f());
                ptr::copy_nonoverlapping(data.as_ptr(), ptr.as_ptr().cast(), data.len());
                ptr
            }
        };
        unsafe { ptr.as_mut() }
    }

    /// 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<&'a mut [T; N]>
    where
        F: FnOnce() -> [T; N],
    {
        let layout = Layout::new::<[T; N]>();
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, true);

        // We can potentially create a jbe instruction directly to the
        // slow path, but to do so when this is inlined, we need to
        // put the parameter in a register, or on the stack.
        // This is slow for larger sizes, so we cap this behaviour at 16 bytes.
        let mut ptr: NonNull<[T; N]> = if layout.size() <= 16 {
            unsafe {
                if take_slow_path {
                    cold_path();
                    self.try_alloc_slow_with::<_, [T; N]>(f)?
                } else {
                    let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                    debug_assert_eq!(new_top as usize % layout.align(), 0);
                    debug_assert!(new_top as usize + layout.size() <= top as usize);
                    let ptr = (new_top.cast::<[T; N]>()).as_mut_unchecked();
                    let data = mem::ManuallyDrop::new(f());
                    ptr::copy_nonoverlapping(data.as_ptr(), ptr.as_mut_ptr(), data.len());
                    self.top.set(new_top);
                    NonNull::new_unchecked(new_top.cast::<[T; N]>())
                }
            }
        } else {
            unsafe {
                let ptr: NonNull<[T; N]> = if take_slow_path {
                    cold_path();
                    self.try_alloc_slow_raw::<[T; N]>()?
                } else {
                    let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                    debug_assert_eq!(new_top as usize % layout.align(), 0);
                    debug_assert!(new_top as usize + layout.size() <= top as usize);
                    self.top.set(new_top);
                    NonNull::new_unchecked(new_top.cast::<[T; N]>())
                };
                let data = mem::ManuallyDrop::new(f());
                ptr::copy_nonoverlapping(data.as_ptr(), ptr.as_ptr().cast(), data.len());
                ptr
            }
        };
        Ok(unsafe { ptr.as_mut() })
    }

    /// 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>) -> &'a mut [T] {
        // This unwrap is safe, because a vector is passed in, so its layout is
        // clearly representable.
        let len = data.len();
        let ptr = self.alloc_array_raw::<T>(len);
        unsafe {
            let data = mem::ManuallyDrop::new(data);
            ptr::copy_nonoverlapping(data.as_ptr(), ptr.cast::<T>(), data.len());
        }
        unsafe { core::slice::from_raw_parts_mut(ptr, len) }
    }

    /// Fallible version of [`alloc_boxed_slice`](Self::alloc_boxed_slice).
    #[inline]
    pub fn try_alloc_vec<T>(&self, data: Vec<T>) -> Result<&'a mut [T]> {
        // This unwrap is safe, because a vector is passed in, so its layout is
        // clearly representable.
        let len = data.len();
        let ptr = self.try_alloc_array_raw::<T>(len)?;
        unsafe {
            let data = mem::ManuallyDrop::new(data);
            ptr::copy_nonoverlapping(data.as_ptr(), ptr.cast::<T>(), data.len());
        }
        Ok(unsafe { core::slice::from_raw_parts_mut(ptr, len) })
    }

    /// 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]>) -> &'a mut [T] {
        self.alloc_vec(data.into_vec())
    }

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

    /// 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) -> &'a mut [T; N]
    where
        F: FnMut(usize) -> T,
    {
        let layout = Layout::new::<[T; N]>();
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, true);

        // We can potentially create a jbe instruction directly to the
        // slow path, but to do so when this is inlined, we need to
        // put the parameter in a register, or on the stack.
        // This is slow for larger sizes, so we cap this behaviour at 16 bytes.
        let mut ptr: NonNull<[T; N]> = if layout.size() <= 16 {
            unsafe {
                if take_slow_path {
                    cold_path();
                    self.alloc_slow_with::<_, [T; N]>(|| array::from_fn(f))
                } else {
                    let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                    debug_assert_eq!(new_top as usize % layout.align(), 0);
                    debug_assert!(new_top as usize + layout.size() <= top as usize);
                    let p = new_top.cast::<T>();
                    Self::fill_with_fn(p, N, f);
                    self.top.set(new_top);
                    NonNull::new_unchecked(new_top.cast::<[T; N]>())
                }
            }
        } else {
            unsafe {
                let ptr: NonNull<[T; N]> = if take_slow_path {
                    cold_path();
                    self.alloc_layout_slow(layout).cast()
                } else {
                    let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                    debug_assert_eq!(new_top as usize % layout.align(), 0);
                    debug_assert!(new_top as usize + layout.size() <= top as usize);
                    self.top.set(new_top);
                    NonNull::new_unchecked(new_top.cast::<[T; N]>())
                };
                Self::fill_with_fn(ptr.as_ptr().cast(), N, f);
                ptr.cast::<[T; N]>()
            }
        };
        unsafe { ptr.as_mut() }
    }

    /// Allocate an array of known size, populated element-by-element.
    #[inline]
    pub fn try_alloc_sized_slice_fill_with<T, const N: usize, F>(
        &self,
        f: F,
    ) -> Result<&'a mut [T; N]>
    where
        F: FnMut(usize) -> T,
    {
        let layout = Layout::new::<[T; N]>();
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, true);

        // We can potentially create a jbe instruction directly to the
        // slow path, but to do so when this is inlined, we need to
        // put the parameter in a register, or on the stack.
        // This is slow for larger sizes, so we cap this behaviour at 16 bytes.
        let mut ptr: NonNull<[T; N]> = if layout.size() <= 16 {
            unsafe {
                if take_slow_path {
                    cold_path();
                    self.try_alloc_slow_with::<_, [T; N]>(|| array::from_fn(f))?
                } else {
                    let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                    debug_assert_eq!(new_top as usize % layout.align(), 0);
                    debug_assert!(new_top as usize + layout.size() <= top as usize);
                    let p = new_top.cast::<T>();
                    Self::fill_with_fn(p, N, f);
                    self.top.set(new_top);
                    NonNull::new_unchecked(new_top.cast::<[T; N]>())
                }
            }
        } else {
            unsafe {
                let ptr: NonNull<[T; N]> = if take_slow_path {
                    cold_path();
                    self.try_alloc_layout_slow(layout)?.cast()
                } else {
                    let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                    debug_assert_eq!(new_top as usize % layout.align(), 0);
                    debug_assert!(new_top as usize + layout.size() <= top as usize);
                    self.top.set(new_top);
                    NonNull::new_unchecked(new_top.cast::<[T; N]>())
                };
                Self::fill_with_fn(ptr.as_ptr().cast(), N, f);
                ptr.cast::<[T; N]>()
            }
        };
        Ok(unsafe { ptr.as_mut() })
    }

    /// 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) -> &'a mut [T]
    where
        F: FnMut(usize) -> T,
    {
        let ptr = self.alloc_array_raw(len);
        unsafe {
            Self::fill_with_fn(ptr, len, f);
            core::slice::from_raw_parts_mut(ptr, len)
        }
    }

    /// 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<&'a mut [T]>
    where
        F: FnMut(usize) -> T,
    {
        let ptr = self.try_alloc_array_raw(len)?;
        unsafe {
            Self::fill_with_fn(ptr, len, f);
            Ok(core::slice::from_raw_parts_mut(ptr, len))
        }
    }

    /// Allocate a slice, populating all slots uniformly by a fallible function.
    /// Returns `Err(e)` if `f` returns `Err`, with no partially-allocated state.
    ///
    /// # 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,
        mut f: F,
    ) -> core::result::Result<&'a mut [T], E>
    where
        F: FnMut(usize) -> core::result::Result<T, E>,
    {
        let layout = Layout::array::<T>(len).unwrap();
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, false);

        unsafe {
            let ptr: *mut T = if take_slow_path {
                cold_path();
                self.alloc_layout_slow(layout).cast().as_ptr()
            } else {
                let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                debug_assert_eq!(new_top as usize % layout.align(), 0);
                debug_assert!(new_top as usize + layout.size() <= top as usize);
                self.top.set(new_top);
                new_top.cast::<T>()
            };
            let mut p = ptr;
            // Optimizes to a memcpy, when source is a contiguous array,
            // and is large, or of unknown size.
            for i in 0..len {
                match f(i) {
                    Ok(a) => {
                        ptr::write(p, a);
                        p = p.add(1);
                    }
                    Err(e) => {
                        cold_path();
                        if take_slow_path {
                            cold_path();
                            let header = self.header().as_ref_unchecked();
                            self.top
                                .set(self.bottom.get().byte_add(header.capacity).cast::<u8>());
                        } else {
                            // Fast path: restore original top to undo the partial alloc.
                            self.top.set(top);
                        }
                        return Err(e);
                    }
                }
            }
            Ok(core::slice::from_raw_parts_mut(ptr, len))
        }
    }

    /// 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]) -> &'a mut [T; N] {
        self.alloc_sized_slice_with(|| *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<&'a mut [T; N]> {
        self.try_alloc_sized_slice_with(|| *data)
    }

    /// Clone a fixed-size array into the arena. Returns a mutable reference to the copy.
    #[inline]
    pub fn alloc_sized_slice_clone<T: Clone, const N: usize>(
        &self,
        data: &[T; N],
    ) -> &'a mut [T; N] {
        self.alloc_sized_slice_fill_with(|n| data[n].clone())
    }

    /// 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<&'a mut [T; N]> {
        self.try_alloc_sized_slice_fill_with(|n| data[n].clone())
    }

    /// Allocate a fixed-size array filled with [`Default`] values.
    #[inline]
    pub fn alloc_sized_slice_fill_default<T: Default, const N: usize>(&self) -> &'a mut [T; N] {
        self.alloc_sized_slice_fill_with(|_| Default::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<&'a mut [T; N]> {
        self.try_alloc_sized_slice_fill_with(|_| Default::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) -> &'a mut [T] {
        self.alloc_slice_fill_with(len, |_| Default::default())
    }

    /// 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<&'a mut [T]> {
        self.try_alloc_slice_fill_with(len, |_| Default::default())
    }

    /// Allocate a slice from an iterator with 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) -> &'a mut [T]
    where
        I: IntoIterator<Item = T>,
        <I as IntoIterator>::IntoIter: ExactSizeIterator,
    {
        let iter = iter.into_iter();
        let len = iter.len();
        let ptr = self.alloc_array_raw(len);
        unsafe {
            Self::fill_with_iter_exact(ptr, len, iter);
            core::slice::from_raw_parts_mut(ptr, len)
        }
    }

    /// 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<&'a mut [T]>
    where
        I: IntoIterator<Item = T>,
        <I as IntoIterator>::IntoIter: ExactSizeIterator,
    {
        let iter = iter.into_iter();
        let len = iter.len();
        let ptr = self.try_alloc_array_raw(len)?;
        unsafe {
            Self::fill_with_iter_exact(ptr, len, iter);
            Ok(core::slice::from_raw_parts_mut(ptr, len))
        }
    }

    /// Allocate a slice from an iterator with trusted known size.
    /// Returns a mutable reference to the allocated slice.
    ///
    /// # Panics
    ///
    /// 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) -> &'a mut [T]
    where
        I: IntoIterator<Item = T>,
        <I as IntoIterator>::IntoIter: TrustedLen,
    {
        let iter = iter.into_iter();
        let (len, _) = iter.size_hint();
        let ptr = self.alloc_array_raw(len);
        unsafe {
            Self::fill_with_iter_trusted(ptr, iter);
            core::slice::from_raw_parts_mut(ptr, len)
        }
    }

    /// Copy a string into the arena. Returns a mutable reference to the copy.
    #[inline]
    pub fn alloc_str(&self, s: &str) -> &'a mut str {
        let slice = self.alloc_slice_copy(s.as_bytes());
        unsafe { core::str::from_utf8_unchecked_mut(slice) }
    }

    /// Fallible version of [`alloc_str`](Self::alloc_str).
    #[inline]
    pub fn try_alloc_str(&self, s: &str) -> Result<&'a mut str> {
        let slice = self.try_alloc_slice_copy(s.as_bytes())?;
        Ok(unsafe { core::str::from_utf8_unchecked_mut(slice) })
    }

    /// 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 {
        let header = unsafe { self.header().as_ref_unchecked() };
        let mut res = header.capacity;
        res += count_bytes_allocated_in_chunk_seq(header.next);
        res += count_bytes_allocated_in_chunk_seq(self.fresh_chunks.get());
        res
    }

    /// 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 {
        let header = unsafe { self.header().as_ref_unchecked() };
        let mut res = header.capacity;
        res += count_bytes_allocated_in_chunk_seq_with_meta(header.next);
        res += count_bytes_allocated_in_chunk_seq_with_meta(self.fresh_chunks.get());
        res
    }

    /// 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` — the borrow checker prevents using the outer
    /// `ArenaRef` directly inside the closure.
    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.with_scope_unchecked(f) }
    }

    /// Same as `with_scope` but takes `&mut self`.  Used by `Arena::with_scope`
    /// which has no lifetime parameter and thus cannot prevent holding alloc
    /// references across the call.
    pub(crate) unsafe fn with_scope_unchecked<'brand, F, R>(&'brand mut self, f: F) -> R
    where
        F: for<'scope> FnOnce(&'scope mut ArenaRef<'scope>) -> R,
    {
        // SAFETY: checkpoint saves Cell fields, restore rewinds — both safe
        // because Cells provide interior mutability through &self.
        let checkpoint = unsafe { self.checkpoint() };
        // Use a raw pointer to avoid tying the reborrow to 'a (ArenaRef is
        // invariant in 'a, so f(self) would borrow for 'a and prevent the
        // subsequent self.restore_checkpoint).
        //
        // Safety: 'scope is universally quantified (HRTB), so references
        // returned by the scoped arena are bounded by the closure's scope.
        let inner = core::ptr::from_mut(self);
        let result = f(unsafe { &mut *inner });
        unsafe { self.restore_checkpoint(checkpoint) };
        result
    }

    /// Restores the arena to the state at which the given checkpoint was saved.
    unsafe fn restore_checkpoint(&self, checkpoint: Checkpoint) {
        let mut header = self.header();
        let mut fresh_head = self.fresh_chunks.get();

        // Walk the chunk chain from the current header back toward the
        // checkpoint.  Move each visited chunk to the fresh-chunks list
        // (they will be reused, not freed).
        //
        // The FIRST_HEADER sentinel (capacity 0) is never linked into the
        // chain, so we may hit null before finding the checkpoint header.
        // That's fine — the sentinel is the implicit root.
        while header != checkpoint.header && !header.is_null() {
            // Read/write chunk header fields through raw pointers to avoid
            // creating &mut references that would retag and attach stale
            // borrow tags to the stored pointers.
            let next_header = unsafe { (*header).next };
            unsafe { ptr::write(&raw mut (*header).next, fresh_head) };
            fresh_head = header;
            header = next_header;
        }

        self.top.set(checkpoint.top);
        let new_bottom = unsafe { checkpoint.header.add(1) };
        self.bottom.set(new_bottom);
        self.fresh_chunks.set(fresh_head);
    }

    /// Returns a checkpoint with which the arena can be
    /// reverted to a previous state.
    /// Note that no chunks are actually freed, they're just put
    /// into the reuse list. To free these chunks, call
    /// [`free_unused_chunks`](Arena::free_unused_chunks).
    unsafe fn checkpoint(&self) -> Checkpoint {
        Checkpoint {
            top: self.top.get(),
            header: self.header(),
        }
    }

    /// Copy a string literal into the arena. Returns a mutable reference to the copy.
    /// This is faster than [`alloc_str`](ArenaRef::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`](ArenaRef::alloc_str).
    #[inline]
    pub fn alloc_str_lit(&self, s: &'static str) -> &'a mut str {
        let slice = self.alloc_slice_lit_copy(s.as_bytes());
        unsafe { core::str::from_utf8_unchecked_mut(slice) }
    }

    /// Fallible version of [`alloc_str_lit`](Self::alloc_str_lit).
    #[inline]
    pub fn try_alloc_str_lit(&self, s: &'static str) -> Result<&'a mut str> {
        let slice = self.try_alloc_slice_lit_copy(s.as_bytes())?;
        Ok(unsafe { core::str::from_utf8_unchecked_mut(slice) })
    }
}

/// Opaque snapshot of arena state, produced by [`ArenaRef::checkpoint`].
struct Checkpoint {
    top: *mut u8,
    header: *mut ChunkHeader,
}

// --- Test helpers (doc-hidden) ---

impl ArenaRef<'_> {
    /// Returns whether the fresh-chunks list is non-empty.
    #[doc(hidden)]
    pub fn has_fresh_chunks(&self) -> bool {
        !self.fresh_chunks.get().is_null()
    }

    /// Returns the number of free bytes remaining in the current chunk.
    #[doc(hidden)]
    pub fn chunk_capacity(&self) -> usize {
        self.top.get() as usize - self.bottom.get() as usize
    }
}

// --- Private functions ---

impl ArenaRef<'_> {
    /// This wrapper is never inlined, which keeps the amount of code generated
    /// inline as the fast path minimal.
    /// It immediately calls out to the non-generic `alloc_layout_slow`, which
    /// keeps the overall emitted code size small.
    #[inline(never)]
    fn alloc_slow_with<F, T>(&self, f: F) -> NonNull<T>
    where
        F: FnOnce() -> T,
    {
        let p = self.alloc_slow_raw::<T>();
        unsafe { ptr::write(p.as_ptr(), f()) };
        p
    }

    #[inline(never)]
    fn try_alloc_slow_with<F, T>(&self, f: F) -> Result<NonNull<T>>
    where
        F: FnOnce() -> T,
    {
        let p = self.try_alloc_slow_raw::<T>()?;
        unsafe { ptr::write(p.as_ptr(), f()) };
        Ok(p)
    }

    /// Fill a buffer with elements from an iterator.
    /// The buffer must be at least the size of the iterator, otherwise, this function
    /// you'll get buffer overflows.
    unsafe fn fill_with_iter_exact<T, I: ExactSizeIterator<Item = T>>(
        mut p: *mut T,
        len: usize,
        mut iter: I,
    ) {
        for _ in 0..len {
            let item = iter
                .next()
                .expect("Iterator didn't contain enough elements");
            unsafe {
                ptr::write(p, item);
                p = p.add(1);
            }
        }
    }

    /// Fill a buffer with elements from an iterator.
    /// The buffer must be at least the size of the iterator, otherwise, this function
    /// you'll get buffer overflows.
    #[cfg(nightly)]
    unsafe fn fill_with_iter_trusted<T, I: TrustedLen<Item = T>>(mut p: *mut T, iter: I) {
        for next in iter {
            unsafe {
                ptr::write(p, next);
                p = p.add(1);
            }
        }
    }

    /// Fill a buffer with elements from a function.
    unsafe fn fill_with_fn<T, F: FnMut(usize) -> T>(mut p: *mut T, len: usize, mut f: F) {
        for i in 0..len {
            unsafe {
                ptr::write(p, f(i));
                p = p.add(1);
            }
        }
    }

    /// 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]
    fn alloc_array_raw<T>(&self, len: usize) -> *mut T {
        let layout = Layout::array::<T>(len).unwrap();
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, false);

        unsafe {
            if take_slow_path {
                cold_path();
                self.alloc_layout_slow(layout).cast().as_ptr()
            } else {
                let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                debug_assert_eq!(new_top as usize % layout.align(), 0);
                debug_assert!(new_top as usize + layout.size() <= top as usize);
                self.top.set(new_top);
                new_top.cast::<T>()
            }
        }
    }

    /// Allocate space for an object with the given [`Layout`].
    ///
    /// The returned pointer points at uninitialized memory.
    #[inline]
    fn try_alloc_array_raw<T>(&self, len: usize) -> Result<*mut T> {
        let layout = Layout::array::<T>(len).unwrap();
        let top = self.top.get();
        let bottom = self.bottom.get() as usize;

        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        let take_slow_path = Self::need_slow_path(layout, top, bottom, false);

        Ok(unsafe {
            if take_slow_path {
                cold_path();
                self.try_alloc_layout_slow(layout)?.cast().as_ptr()
            } else {
                let new_top = top.byte_sub(extra_bytes_needed).byte_sub(layout.size());
                debug_assert_eq!(new_top as usize % layout.align(), 0);
                debug_assert!(new_top as usize + layout.size() <= top as usize);
                self.top.set(new_top);
                new_top.cast::<T>()
            }
        })
    }

    // Doesn't write to the allocated memory
    fn alloc_slow_raw<T>(&self) -> NonNull<T> {
        let layout = Layout::new::<T>();
        self.alloc_layout_slow(layout).cast()
    }

    fn try_alloc_slow_raw<T>(&self) -> Result<NonNull<T>> {
        let layout = Layout::new::<T>();
        Ok(self.try_alloc_layout_slow(layout)?.cast())
    }

    /// Detect whether to take the slow path.
    /// The `comptime_layout` parameter is a compile-time constant at all call
    /// sites, so the compiler eliminates one side of the outer guard after
    /// inlining.
    #[inline(always)]
    fn need_slow_path(layout: Layout, top: *mut u8, bottom: usize, comptime_layout: bool) -> bool {
        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());

        // The most padding you'll ever need is the alignment - 1.
        let max_padding = layout.align() - 1;

        // The maximum space that an allocation can take up is its size
        // plus the maximum amount of padding added to align the top.
        let max_space: usize = layout.size() + max_padding;
        let top = top as usize;

        // Both halves of this condition are compile-time constants:
        if comptime_layout && max_space < SAFE_SIZE {
            // No overflow possible.
            if max_space <= MIN_TOP {
                // Optimal path: subtract alignment + size from top, unified comparison.
                let top = top - extra_bytes_needed;
                let top = top - layout.size();
                bottom > top
            } else if max_padding < MIN_TOP {
                let top = top - extra_bytes_needed;
                bottom + layout.size() > top
            } else {
                bottom + layout.size() + extra_bytes_needed > top
            }
        } else if max_padding
            // Layout::size can never be greater than isize::MAX
            .checked_add(isize::MAX as usize)
            .is_some_and(|a| a < SAFE_SIZE)
        {
            // This codepath is one more instruction, but one less conditional
            // jump, and seems to be worth it.
            if max_padding < MIN_TOP {
                let top = top - extra_bytes_needed;
                bottom + layout.size() > top
            } else {
                bottom + extra_bytes_needed + layout.size() > top
            }
        } else {
            // Even when we don't have a size available at compile-time,
            // we should still have the alignment...
            if max_padding < MIN_TOP {
                let top = top - extra_bytes_needed;
                top.checked_sub(layout.size()).is_none_or(|t| t < bottom)
            } else {
                Self::need_slow_path_checked(layout, top as *mut u8, bottom)
            }
        }
    }

    /// Detect whether to take the slow path.
    /// This is a slow version which always uses checked arithmetic.
    #[inline(always)]
    fn need_slow_path_checked(layout: Layout, top: *mut u8, bottom: usize) -> bool {
        let extra_bytes_needed = ArenaRef::extra_bytes_needed(top, layout.align());
        (top as usize)
            .checked_sub(extra_bytes_needed)
            .and_then(|p| p.checked_sub(layout.size()))
            .is_none_or(|p| p < bottom)
    }

    #[inline(always)]
    fn extra_bytes_needed(ptr: *const u8, alignment: usize) -> usize {
        assert!(alignment.is_power_of_two());

        // ptr           = 11010
        // 8             = 01000
        // 8 - 1         = 00111
        // ptr & (8 - 1) = 00010

        // Also works if alignment > ptr

        // ptr           = 00010
        // 8             = 01000
        // 8 - 1         = 00111
        // ptr & (8 - 1) = 00010
        ptr as usize & (alignment - 1)
    }

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

    // In a fresh chunk, how many extra bytes do we need to allocate
    // as padding?
    fn extra_fresh_bytes_needed(alignment: usize) -> usize {
        let start: usize = if alignment > HEADER_SIZE {
            alignment - HEADER_SIZE
        } else {
            HEADER_SIZE
        };
        start & (alignment - 1)
    }

    // Slow path: allocate a new chunk (or reuse a fresh one) and return a
    // pointer to the reserved space.
    fn alloc_layout_slow_internal(&self, layout: Layout) -> Result<NonNull<u8>> {
        let extra_fresh_bytes = Self::extra_fresh_bytes_needed(layout.align());
        let min_size = layout.size() + extra_fresh_bytes;

        // Scan the fresh-chunks list for one that's big enough for this allocation.
        // Pop and discard undersized chunks (they can't hold future allocations either,
        // so we deallocate them to avoid leaking).
        loop {
            let fresh = self.fresh_chunks.get();
            if fresh.is_null() {
                break;
            }
            let capacity = unsafe { (*fresh).capacity };
            if capacity >= min_size {
                // This fresh chunk is big enough. Reuse it.
                let next = self.header();
                unsafe {
                    self.fresh_chunks.set((*fresh).next);
                    (*fresh).next = next;
                }
                let new_bottom = unsafe { fresh.add(1) };
                self.bottom.set(new_bottom);
                let top = unsafe { new_bottom.byte_add(capacity).cast::<u8>() };
                let extra = Self::extra_bytes_needed(top, layout.align());
                let new_top = unsafe { top.byte_sub(layout.size() + extra) };
                self.top.set(new_top);
                debug_assert_eq!(new_top as usize % layout.align(), 0);
                debug_assert!(new_top as usize + layout.size() <= top as usize);
                return Ok(unsafe { NonNull::new_unchecked(new_top) });
            }
            // Chunk too small — pop it and deallocate, then check the next one.
            unsafe {
                self.fresh_chunks.set((*fresh).next);
                let layout = Layout::from_size_align_unchecked(capacity + HEADER_SIZE, CHUNK_ALIGN);
                dealloc(fresh.cast::<u8>(), layout);
            }
        }

        // No usable fresh chunk — allocate a new one.
        let next = self.header();
        // Don't link the FIRST_HEADER sentinel into the chain
        let next = if unsafe { (*next).capacity == 0 } {
            ptr::null_mut()
        } else {
            next
        };
        let size = self.next_chunk_size().max(min_size);
        let new_chunk = ChunkHeader::try_new(size, next)?;
        let new_bottom = unsafe { new_chunk.add(1) };
        self.bottom.set(new_bottom);
        let top = unsafe { new_bottom.byte_add(size).cast::<u8>() };
        let extra = Self::extra_bytes_needed(top, layout.align());
        let new_top = unsafe { top.byte_sub(layout.size() + extra) };
        self.top.set(new_top);
        debug_assert_eq!(new_top as usize % layout.align(), 0);
        debug_assert!(new_top as usize + layout.size() <= top as usize);
        Ok(unsafe { NonNull::new_unchecked(new_top) })
    }

    #[inline(never)]
    fn try_alloc_layout_slow(&self, layout: Layout) -> Result<NonNull<u8>> {
        self.alloc_layout_slow_internal(layout)
    }

    #[inline(never)]
    fn alloc_layout_slow(&self, layout: Layout) -> NonNull<u8> {
        match self.alloc_layout_slow_internal(layout) {
            Ok(a) => a,
            Err(e) => {
                if matches!(e, Error::OOM) {
                    core_alloc::alloc::handle_alloc_error(layout);
                }
                panic!("{}", e)
            }
        }
    }

    fn next_chunk_size(&self) -> usize {
        let header = self.header();
        let cap = unsafe { header.as_ref_unchecked().capacity };
        let res = (cap + HEADER_SIZE).saturating_mul(2) - HEADER_SIZE;
        res.max(INITIAL_CHUNK_CAPACITY)
    }
}

impl Drop for ArenaRef<'_> {
    fn drop(&mut self) {
        // Merge state down to the base Arena so forwarding methods see it.
        let arena = unsafe { &*self.arena_ptr };
        arena.top.set(self.top.get());
        arena.bottom.set(self.bottom.get());
        arena.fresh_chunks.set(self.fresh_chunks.get());
    }
}
