bumpish 0.4.1

use crate::util::{self, ALLOC_OVERHEAD, TypeProps};
use alloc::alloc::{Layout, alloc, dealloc, handle_alloc_error};
use core::cell::Cell;
use core::marker::PhantomData;
use core::mem::{MaybeUninit, offset_of};
use core::ptr::{self, NonNull};

struct ChunkHeader {
    /// The offset of the finger of the data buffer relative to the start of the
    /// buffer, i.e. the `data` field.
    ///
    /// This offset must always be less than or equal to `end`.
    ptr: Cell<usize>,

    /// The offset of the end of the data buffer relative to the start of the
    /// buffer, i.e. the `data` field.
    end: usize,

    /// Link to the previous chunk.
    ///
    /// Note that the last node in the `prev` linked list is the dead chunk,
    /// whose `prev` link points to itself.
    prev: Cell<NonNull<u8>>,

    /// Link to the next chunk.
    ///
    /// Note that the last node in the `next` linked list is the dead chunk,
    /// whose `next` link points to itself.
    next: Cell<NonNull<u8>>,
}

pub(super) struct ChunkProps<T> {
    /// Size of the chunk of memory in bytes.
    size: usize,
    /// Number of elements that the chunk can hold.
    capacity: usize,

    phantom: PhantomData<T>,
}

impl<T> ChunkProps<T> {
    /// Alignment of the chunk.
    const ALIGN: usize = {
        assert!(Chunk::<T, 0>::ALIGN == Chunk::<T, 1>::ALIGN);
        Chunk::<T, 1>::ALIGN
    };

    /// Offset of the `data` field from the start of the chunk.
    const START_OFFSET: usize = offset_of!(Chunk<T, 1>, data);

    /// Recommended first chunk size.
    pub(super) const FIRST_SIZE: usize = {
        let chunk_512b = 0x200;
        let chunk_2elem = (2 * T::SIZE + Self::START_OFFSET).next_power_of_two();
        let rough_chunk_size = util::max(chunk_512b, chunk_2elem);
        Self::from_pow2(rough_chunk_size).size
    };

    /// Chunk size enough for at least one element.
    pub(super) const MIN_SIZE: usize =
        Self::from_pow2(Self::min_size_for(1).next_power_of_two()).size;

    /// Maximal possible chunk size. Is equal to 4 GiB, if address space is big
    /// enough. Otherwise, it is equal to `size of address space size / 4`.
    const MAX_POW2: usize = {
        if usize::BITS >= 64 {
            4 << 30 // 4 GiB
        } else {
            (usize::MAX >> 2) + 1
        }
    };

    pub(super) const fn min_size_for(num_of_elems: usize) -> usize {
        num_of_elems * T::SIZE + Self::START_OFFSET + ALLOC_OVERHEAD
    }

    pub(super) const fn from_pow2(pow_of_two: usize) -> Self {
        debug_assert!(pow_of_two.is_power_of_two());
        if const { T::IS_ZST } {
            Self {
                size: 0,
                capacity: 0,
                phantom: PhantomData,
            }
        } else {
            let mut size = pow_of_two - ALLOC_OVERHEAD;
            size -= Self::START_OFFSET;
            let capacity = size - (size % T::SIZE);
            size = capacity + Self::START_OFFSET;
            Self {
                size,
                capacity,
                phantom: PhantomData,
            }
        }
    }
}

#[repr(C)]
pub(super) struct Chunk<T, const N: usize> {
    header: ChunkHeader,
    data: [MaybeUninit<T>; N],
}

impl<T, const N: usize> Chunk<T, N> {
    pub(super) const fn new() -> Self {
        Self {
            header: ChunkHeader {
                ptr: Cell::new(ChunkProps::<T>::START_OFFSET),
                end: ChunkProps::<T>::START_OFFSET + N * T::SIZE,
                prev: Cell::new(DEAD_CHUNK.get()),
                next: Cell::new(DEAD_CHUNK.get()),
            },
            data: [const { MaybeUninit::uninit() }; N],
        }
    }

    pub(super) const fn get(&mut self) -> ChunkRef<T> {
        let ptr = NonNull::from_mut(self);
        ChunkRef(ptr.cast(), PhantomData)
    }

    /// First heap allocated chunk
    pub(super) fn next_chunk(&self) -> ChunkRef<T> {
        unsafe { ChunkRef::from(self.header.next.get()) }
    }

    pub(super) unsafe fn set_next_chunk(&mut self, chunk: ChunkRef<T>) {
        self.header.next.set(chunk.0);
    }

    /// Current heap allocated chunk
    #[inline(always)]
    pub(super) fn current_chunk(&self) -> ChunkRef<T> {
        unsafe { ChunkRef::from(self.header.prev.get()) }
    }

    pub(super) unsafe fn set_current_chunk(&mut self, chunk: ChunkRef<T>) {
        self.header.prev.set(chunk.0);
    }
}

/// A wrapper around a pointer to a chunk of data.
///
/// # Safety
///
/// Almost every operation on this type is unsafe, as it assumes the pointer
/// points to a valid chunk of memory, what is not guaranteed because of it's
/// not restricted by a proper lifetime.
pub(super) struct ChunkRef<T>(NonNull<u8>, PhantomData<*mut T>);

impl<T> ChunkRef<T> {
    pub(super) const unsafe fn from(ptr: NonNull<u8>) -> Self {
        ChunkRef(ptr.cast(), PhantomData)
    }
    /// Allocates a new chunk with a size based on the next power of two of the
    /// rough size.
    ///
    /// Returns a tuple of the new chunk ref and the capacity of the chunk.
    pub(super) fn new(rough_size: usize) -> (Self, usize) {
        let mut pow2 = rough_size.next_power_of_two();
        if pow2 > ChunkProps::<T>::MAX_POW2 {
            pow2 = ChunkProps::<T>::MAX_POW2;
        }
        let mut props = ChunkProps::<T>::from_pow2(pow2);
        let chunk_align = ChunkProps::<T>::ALIGN;

        let chunk_ptr = loop {
            let chunk_layout =
                unsafe { Layout::from_size_align_unchecked(props.size, chunk_align) };

            let chunk_ptr = unsafe { alloc(chunk_layout) };
            if !chunk_ptr.is_null() {
                debug_assert!(util::ptr_is_aligned_to(chunk_ptr, ChunkProps::<T>::ALIGN));
                break chunk_ptr;
            }

            // if couldn't get a new chunk, try to shrink the chunk size by half
            pow2 >>= 1;
            props = ChunkProps::<T>::from_pow2(pow2);

            if props.size < ChunkProps::<T>::MIN_SIZE {
                handle_alloc_error(chunk_layout);
            }
        };

        unsafe {
            let chunk_ptr = NonNull::new_unchecked(chunk_ptr);
            let header_ptr = chunk_ptr.cast::<ChunkHeader>();
            header_ptr.write(ChunkHeader {
                ptr: Cell::new(ChunkProps::<T>::START_OFFSET),
                end: props.size,
                prev: Cell::new(DEAD_CHUNK.get()),
                next: Cell::new(DEAD_CHUNK.get()),
            });

            let chunk = Self(chunk_ptr, PhantomData);
            (chunk, props.capacity)
        }
    }

    /// Returns a chunk reference to the `DEAD_CHUNK` static instance.
    pub(super) fn dead() -> Self {
        Self(DEAD_CHUNK.get(), PhantomData)
    }

    /// Drops elements from the chunk without deallocating the chunk itself.
    pub(super) unsafe fn drop_elements(&self) {
        unsafe {
            if const { core::mem::needs_drop::<T>() } {
                let ptr = self.start();
                let len = self.len();
                // TODO: replace with `slice::from_mut_ptr_range` when it is
                // stable, to avoid division in `ChunkRef::len`
                let slice = ptr::slice_from_raw_parts_mut(ptr.as_ptr(), len);
                ptr::drop_in_place(slice);
            }
            self.reset_ptr()
        }
    }

    /// Drops all elements in the chunk and deallocates it.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the chunk is no longer in use.
    pub(super) unsafe fn drop(self) {
        unsafe {
            self.drop_elements();
            let layout = Layout::from_size_align_unchecked(self.size(), ChunkProps::<T>::ALIGN);
            dealloc(self.get().as_ptr(), layout);
        }
    }

    pub(super) fn is(&self, other: Self) -> bool {
        ptr::eq(self.0.as_ptr(), other.0.as_ptr())
    }

    /// Returns `true` if the chunk is the `DEAD_CHUNK`.
    pub(super) unsafe fn is_dead(&self) -> bool {
        ptr::eq(self.0.cast().as_ptr(), &DEAD_CHUNK.0)
    }

    /// Returns the total capacity of the chunk in bytes.
    pub(super) unsafe fn capacity(&self) -> usize {
        let header = unsafe { self.header() };
        header.end - ChunkProps::<T>::START_OFFSET
    }

    /// Returns the number of elements currently stored in the chunk.
    pub(super) const unsafe fn len(&self) -> usize {
        let header = unsafe { self.header() };
        (header.ptr.get() - ChunkProps::<T>::START_OFFSET) / T::SIZE
    }

    /// Returns the size of the chunk.
    pub(super) const unsafe fn size(&self) -> usize {
        unsafe { self.header() }.end
    }

    /// Returns the size of the of pushed elements in bytes.
    pub(super) const unsafe fn data_size(&self) -> usize {
        unsafe { self.header() }.ptr.get() - ChunkProps::<T>::START_OFFSET
    }

    /// Checks if the chunk is full.
    #[inline(always)]
    pub(super) unsafe fn is_full(&self) -> bool {
        let header = unsafe { self.header() };
        header.ptr.get() == header.end
    }

    /// Checks if the chunk does not contain any elements.
    #[inline(always)]
    pub(super) unsafe fn is_empty(&self) -> bool {
        let header = unsafe { self.header() };
        header.ptr.get() == ChunkProps::<T>::START_OFFSET
    }

    /// Returns a pointer to the first element in the chunk.
    ///
    /// # Safety
    ///
    /// The chunk must not be empty.
    pub(super) unsafe fn first_unchecked(&self) -> NonNull<T> {
        debug_assert!(!T::IS_ZST);
        unsafe { self.0.add(ChunkProps::<T>::START_OFFSET).cast() }
    }

    /// Returns a pointer to the last element in the chunk.
    ///
    /// # Safety
    ///
    /// The chunk must not be empty.
    pub(super) unsafe fn last_unchecked(&self) -> NonNull<T> {
        debug_assert!(!T::IS_ZST);
        unsafe { self.0.add(self.header().ptr.get() - T::SIZE).cast() }
    }

    pub(super) unsafe fn next(&self) -> Self {
        unsafe { ChunkRef::from(self.header().next.get()) }
    }

    pub(super) unsafe fn set_next(&self, next: Self) {
        let header = unsafe { self.header() };
        header.next.set(next.0)
    }

    pub(super) unsafe fn prev(&self) -> Self {
        unsafe { ChunkRef::from(self.header().prev.get()) }
    }

    pub(super) unsafe fn set_prev(&self, next: Self) {
        let header = unsafe { self.header() };
        header.prev.set(next.0)
    }

    pub(super) fn get(&self) -> NonNull<u8> {
        self.0
    }

    /// Allocates memory for a new element in the chunk and return a pointer to
    /// it.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the method is called only for non zero-sized
    /// types.
    #[inline(always)]
    pub(super) unsafe fn alloc_element(&self) -> Option<NonNull<T>> {
        debug_assert!(!T::IS_ZST);

        if unsafe { self.is_full() } {
            return None;
        }

        let header = unsafe { self.header() };
        let elem_ptr = unsafe { self.0.byte_add(header.ptr.get()) };
        header.ptr.update(|ptr| ptr + T::SIZE);

        Some(elem_ptr.cast())
    }

    pub(super) unsafe fn dealloc_element(&self) -> Option<NonNull<T>> {
        debug_assert!(!T::IS_ZST);

        if unsafe { self.is_empty() } {
            return None;
        }

        let header = unsafe { self.header() };
        header.ptr.update(|ptr| ptr - T::SIZE);
        let old_elem_ptr = unsafe { self.0.byte_add(header.ptr.get()) };

        Some(old_elem_ptr.cast())
    }

    /// Returns a pointer to the element at the given offset relative to the
    /// start of the chunk.
    pub(super) unsafe fn ptr_by_offset(&self, offset: usize) -> NonNull<T> {
        debug_assert!(ChunkProps::<T>::START_OFFSET <= offset);
        debug_assert!(offset <= unsafe { self.end_offset() });

        unsafe { self.0.byte_add(offset).cast() }
    }

    /// Returns a shared reference to the chunk header.
    #[inline(always)]
    const unsafe fn header(&self) -> &ChunkHeader {
        unsafe { self.0.cast().as_ref() }
    }

    pub(super) unsafe fn ptr(&self) -> NonNull<T> {
        unsafe { self.0.byte_add(self.header().ptr.get()).cast() }
    }

    pub(super) unsafe fn ptr_offset(&self) -> usize {
        unsafe { self.header().ptr.get() }
    }

    /// Makes the chunk empty without dropping any elements and freeing the
    /// memory.
    pub(super) unsafe fn reset_ptr(&self) {
        unsafe { self.header().ptr.set(ChunkProps::<T>::START_OFFSET) }
    }

    pub(super) unsafe fn start(&self) -> NonNull<T> {
        unsafe { self.0.byte_add(ChunkProps::<T>::START_OFFSET).cast() }
    }

    pub(super) unsafe fn start_offset(&self) -> usize {
        ChunkProps::<T>::START_OFFSET
    }

    pub(super) unsafe fn end(&self) -> NonNull<T> {
        unsafe { self.0.byte_add(self.header().end).cast() }
    }

    pub(super) unsafe fn end_offset(&self) -> usize {
        unsafe { self.header().end }
    }

    pub(super) unsafe fn slice_ptr(&self) -> (NonNull<T>, usize) {
        unsafe { (self.start(), self.len()) }
    }
}

impl<T> Copy for ChunkRef<T> {}

impl<T> Clone for ChunkRef<T> {
    fn clone(&self) -> Self {
        *self
    }
}

#[repr(transparent)]
struct DeadChunk(ChunkHeader);

unsafe impl Sync for DeadChunk {}

impl DeadChunk {
    /// Pointer to the header of the chunk.
    const fn get(&'static self) -> NonNull<u8> {
        NonNull::new(self as *const DeadChunk as *mut u8).unwrap()
    }
}

/// Empty chunk that contains only its header. It serves as a sentinel value.
///
/// `ptr` and `end` should be the same.
static DEAD_CHUNK: DeadChunk = DeadChunk(ChunkHeader {
    ptr: Cell::new(size_of::<ChunkHeader>()),
    end: size_of::<ChunkHeader>(),
    prev: Cell::new(DEAD_CHUNK.get()),
    next: Cell::new(DEAD_CHUNK.get()),
});

#[cfg(test)]
mod utest {
    use super::*;

    #[test]
    fn chunk_clone() {
        let mut chunk = Chunk::<usize, 2>::new();
        let chunk_ref = chunk.get();
        #[allow(clippy::clone_on_copy)]
        let chunk_ref_clone = chunk_ref.clone();
        assert!(chunk_ref.is(chunk_ref_clone));
    }

    #[test]
    fn chunkprops_from_pow2() {
        let props = ChunkProps::<()>::from_pow2(2);
        assert_eq!(props.capacity, 0);
        assert_eq!(props.size, 0);
    }
}