bumpish 0.4.0

mod chunk;
mod iter;

pub(crate) use iter::{ChunkIterPtr, IntoIter, IterPtr};

use crate::util::{self, TypeProps};
use chunk::{Chunk, ChunkProps, ChunkRef};
use core::cell::{Cell, UnsafeCell};
use core::ptr::{self, NonNull};

pub(crate) struct RawBump<T, const N: usize> {
    /// The inlined chunk used for small vectors.
    inlined: UnsafeCell<Chunk<T, N>>,

    /// - if `T` is ZST — the total capacity of the vector, including inlined
    ///   capacity, in bytes;
    /// - if `T` is not ZST — number of elements that the vector contains.
    cap_or_len: Cell<usize>,
}

impl<T, const N: usize> RawBump<T, N> {
    pub(crate) const fn new() -> Self {
        Self {
            inlined: UnsafeCell::new(Chunk::new()),
            cap_or_len: if const { T::IS_ZST } {
                Cell::new(0)
            } else {
                Cell::new(N * T::SIZE)
            },
        }
    }

    pub(crate) fn with_capacity(capacity: usize) -> Self {
        let mut core = Self::new();
        if const { !T::IS_ZST } && capacity > N {
            unsafe {
                let rough_size = ChunkProps::<T>::min_size_for(capacity - N);
                let chunk = core.new_chunk(rough_size);
                let inlined = core.inlined.get_mut();
                inlined.set_next_chunk(chunk);
                if const { N == 0 } {
                    inlined.set_current_chunk(chunk);
                }
                debug_assert!(chunk.prev().is_dead());
                debug_assert!(chunk.next().is_dead());
            }
        }
        core
    }

    /// Returns number of ZST elements that the vector holds. Calling this
    /// method for non-ZST types will result in undefined behavior.
    ///
    /// # Safety
    ///
    /// This method shouldn't be called if `T` is not ZST.
    pub(crate) const unsafe fn zst_len(&self) -> usize {
        debug_assert!(T::IS_ZST);
        self.cap_or_len.get()
    }

    /// Returns the capacity of the vector, in number of elements. For ZST
    /// elements, the capacity is always [`usize::MAX`].
    pub(crate) const fn capacity(&self) -> usize {
        if const { T::IS_ZST } {
            usize::MAX
        } else {
            self.cap_or_len.get() / T::SIZE
        }
    }

    /// Returns a pointer to the first element in the vector.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the vector is not emtpy.
    pub(crate) unsafe fn first_unchecked(&self) -> NonNull<T> {
        if const { T::IS_ZST } {
            util::zst_ptr()
        } else {
            let inlined_chunk = unsafe { &mut *self.inlined.get() };
            if const { N == 0 } {
                unsafe { inlined_chunk.next_chunk().first_unchecked() }
            } else {
                unsafe { inlined_chunk.get().first_unchecked() }
            }
        }
    }

    /// Returns a pointer to the last element in the vector.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the vector is not emtpy.
    pub(crate) unsafe fn last_unchecked(&self) -> NonNull<T> {
        if const { T::IS_ZST } {
            util::zst_ptr()
        } else {
            unsafe {
                let inlined_chunk = &mut *self.inlined.get();
                let mut current_chunk = inlined_chunk.current_chunk();
                if current_chunk.is_empty() {
                    current_chunk = current_chunk.prev();
                }
                if const { N > 0 } && current_chunk.is_dead() {
                    current_chunk = inlined_chunk.get();
                }
                current_chunk.last_unchecked()
            }
        }
    }

    /// Returns a pointer to the element at the given index.
    ///
    /// # Safety
    ///
    /// The caller must ensure that `index` is less than the number of elements
    /// in the vector.
    pub(crate) unsafe fn index(&self, index: usize) -> NonNull<T> {
        if const { T::IS_ZST } {
            debug_assert!(index < unsafe { self.zst_len() });
            return util::zst_ptr();
        }
        unsafe {
            let mut offset = index * T::SIZE;
            let mut curr_chunk = if const { N > 0 } {
                ChunkRef::from(self.inlined_ptr().cast())
            } else {
                (*self.inlined.get()).next_chunk()
            };
            loop {
                let data_size = curr_chunk.data_size();
                if offset < data_size {
                    break curr_chunk.start().byte_add(offset);
                } else {
                    offset -= data_size;
                    curr_chunk = curr_chunk.next();
                }
            }
        }
    }

    /// Returns a pointer to the element at the given index, counting from the
    /// back. Indexing starts from `1`, i.e. the last element has index `1`, the
    /// second-to-last has index `2`, and so on.
    ///
    /// # Safety
    ///
    /// The caller must ensure that `index` is less than or equal to the number
    /// of elements in the vector, and greater than `0`.
    pub(crate) unsafe fn index_backward(&self, index: usize) -> NonNull<T> {
        debug_assert!(index > 0);
        if const { T::IS_ZST } {
            debug_assert!(index <= unsafe { self.zst_len() });
            return util::zst_ptr();
        }
        unsafe {
            let mut back_offset = index * T::SIZE;
            let mut curr_chunk = (*self.inlined.get()).current_chunk();
            loop {
                let data_size = curr_chunk.data_size();
                if let Some(offset) = data_size.checked_sub(back_offset) {
                    break curr_chunk.start().byte_add(offset);
                } else {
                    back_offset -= data_size;
                    curr_chunk = curr_chunk.prev();
                    if const { N > 0 } && curr_chunk.is_dead() {
                        let inlined_ptr = self.inlined_ptr();
                        curr_chunk = ChunkRef::from(inlined_ptr.cast());
                    }
                }
            }
        }
    }

    /// Removes the element at `index` by swapping it with the last element and
    /// deallocating the last element.
    ///
    /// # Safety
    ///
    /// `index` must be less than the number of elements in the vector.
    pub(crate) unsafe fn swap_remove(&mut self, index: usize) -> NonNull<T> {
        if const { T::IS_ZST } {
            debug_assert!(index < unsafe { self.zst_len() });
            self.cap_or_len.update(|len| len - 1);
            return util::zst_ptr();
        }
        unsafe {
            let elem_ptr = self.index(index);
            let Some(last_elem_ptr) = self.dealloc_element() else {
                core::hint::unreachable_unchecked();
            };
            elem_ptr.swap(last_elem_ptr);
            last_elem_ptr
        }
    }

    /// Clears the vector, dropping all elements, and deallocating all chunks
    /// except for one that is used as preallocation.
    pub(crate) fn clear(&mut self) {
        unsafe {
            if const { T::IS_ZST } {
                let len = self.zst_len();
                let slice = ptr::slice_from_raw_parts_mut(util::zst_ptr::<T>().as_ptr(), len);
                ptr::drop_in_place(slice);
                self.cap_or_len.set(0);
            } else {
                if const { N > 0 } {
                    self.inlined.get_mut().get().drop_elements();
                }
                let mut first_chunk = self.inlined.get_mut().next_chunk();
                if !first_chunk.is_dead() {
                    loop {
                        let next_chunk = first_chunk.next();
                        if next_chunk.is_dead() {
                            break;
                        }
                        self.drop_chunk(first_chunk);
                        first_chunk = next_chunk;
                    }
                    first_chunk.drop_elements();

                    let inlined_chunk = self.inlined.get_mut();
                    inlined_chunk.set_next_chunk(first_chunk);
                    if const { N > 0 } {
                        inlined_chunk.set_current_chunk(ChunkRef::dead());
                    } else {
                        inlined_chunk.set_current_chunk(first_chunk);
                    }
                    first_chunk.set_prev(ChunkRef::dead());
                }
            }
        }
    }

    /// Pushes a value onto the vector, returning a raw pointer to the allocated
    /// memory.
    #[inline(always)]
    pub(crate) fn allocate(&self) -> NonNull<T> {
        if const { T::IS_ZST } {
            self.cap_or_len.update(|len| {
                len.checked_add(1)
                    .expect("number of ZST elements cannot exceed usize::MAX")
            });
            util::zst_ptr()
        } else {
            unsafe { self.alloc_element() }
        }
    }

    /// Deallocate an element from the vector, returning a pointer at it.
    #[inline]
    pub(crate) fn deallocate(&mut self) -> Option<NonNull<T>> {
        if const { T::IS_ZST } {
            if self.cap_or_len.get() == 0 {
                None
            } else {
                self.cap_or_len.update(|len| len - 1);
                Some(util::zst_ptr())
            }
        } else {
            unsafe { self.dealloc_element() }
        }
    }

    /// Returns an iterator over the elements of the bump allocator.
    pub(crate) fn iter(&self) -> IterPtr<'_, T, N> {
        IterPtr::new(self)
    }

    /// Returns an iterator over the chunks of the bump allocator.
    pub(crate) fn chunk_iter(&self) -> ChunkIterPtr<'_, T, N> {
        ChunkIterPtr::new(self)
    }
}

impl<T, const N: usize> core::iter::IntoIterator for RawBump<T, N> {
    type Item = T;
    type IntoIter = IntoIter<T, N>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter::new(self)
    }
}

impl<T, const N: usize> core::ops::Drop for RawBump<T, N> {
    /// Drops all elements in the bump allocator and deallocates it.
    fn drop(&mut self) {
        self.clear();
        if const { !T::IS_ZST } {
            let current_chunk = self.inlined.get_mut().current_chunk();
            let first_chunk = self.inlined.get_mut().next_chunk();
            unsafe {
                if const { N == 0 } {
                    if !current_chunk.is_dead() {
                        debug_assert!(current_chunk.prev().is_dead());
                        debug_assert!(current_chunk.next().is_dead());
                        debug_assert!(current_chunk.is(first_chunk));
                        self.drop_chunk(current_chunk);
                    } else {
                        debug_assert!(first_chunk.is_dead());
                    }
                } else {
                    //
                    if !first_chunk.is_dead() {
                        debug_assert!(first_chunk.prev().is_dead());
                        debug_assert!(first_chunk.next().is_dead());
                        self.drop_chunk(first_chunk);
                    } else {
                        debug_assert!(current_chunk.is_dead());
                    }
                }
            }
            debug_assert_eq!(self.cap_or_len.get(), N * T::SIZE);
        }
    }
}

impl<T, const N: usize> RawBump<T, N> {
    fn inlined_ptr(&self) -> NonNull<Chunk<T, N>> {
        debug_assert!(!T::IS_ZST);
        unsafe {
            let raw_ptr = self.inlined.get();
            NonNull::new_unchecked(raw_ptr)
        }
    }

    #[inline(always)]
    unsafe fn alloc_element(&self) -> NonNull<T> {
        debug_assert!(!T::IS_ZST);

        let inlined_chunk = unsafe { &mut *self.inlined.get() };
        let current_chunk = inlined_chunk.current_chunk();
        if let Some(ptr) = unsafe { current_chunk.alloc_element() } {
            ptr
        } else {
            if const { N > 0 } {
                let inlined_chunk = unsafe { (&mut *self.inlined.get()).get() };
                if let Some(ptr) = unsafe { inlined_chunk.alloc_element() } {
                    return ptr;
                }
            }
            unsafe { self.alloc_element_slow() }
        }
    }

    #[inline(never)]
    unsafe fn alloc_element_slow(&self) -> NonNull<T> {
        debug_assert!(!T::IS_ZST);

        unsafe {
            let inlined_chunk = &mut *self.inlined.get();
            let first_chunk = inlined_chunk.next_chunk();
            let mut current_chunk = inlined_chunk.current_chunk();
            if current_chunk.is_dead() {
                if const { N == 0 } || first_chunk.is_dead() {
                    debug_assert!(first_chunk.is_dead());

                    let new_chunk = self.new_chunk(ChunkProps::<T>::FIRST_SIZE);
                    current_chunk = new_chunk;
                    inlined_chunk.set_next_chunk(new_chunk);
                } else {
                    current_chunk = first_chunk;
                }
            } else {
                debug_assert!(current_chunk.is_full());
                debug_assert!(!first_chunk.is_dead());

                if current_chunk.next().is_dead() {
                    let rough_size = ChunkProps::<T>::FIRST_SIZE.max(current_chunk.size() << 1);
                    let new_chunk = self.new_chunk(rough_size);
                    current_chunk.set_next(new_chunk);
                    new_chunk.set_prev(current_chunk);
                    current_chunk = new_chunk;
                } else {
                    current_chunk = current_chunk.next();
                    debug_assert!(current_chunk.is_empty());
                }
            }

            inlined_chunk.set_current_chunk(current_chunk);
            current_chunk.alloc_element().unwrap_unchecked()
        }
    }

    #[inline(always)]
    unsafe fn dealloc_element(&mut self) -> Option<NonNull<T>> {
        debug_assert!(!T::IS_ZST);

        let inlined_chunk = self.inlined.get_mut();
        let current_chunk = inlined_chunk.current_chunk();
        if let Some(ptr) = unsafe { current_chunk.dealloc_element() } {
            Some(ptr)
        } else {
            unsafe { self.dealloc_element_slow() }
        }
    }

    #[inline(never)]
    unsafe fn dealloc_element_slow(&mut self) -> Option<NonNull<T>> {
        debug_assert!(!T::IS_ZST);

        unsafe {
            let inlined_chunk = self.inlined.get_mut();
            let mut current_chunk = inlined_chunk.current_chunk();

            if current_chunk.is_dead() {
                if const { N > 0 } {
                    inlined_chunk.get().dealloc_element()
                } else {
                    debug_assert!(inlined_chunk.next_chunk().is_dead());
                    None
                }
            } else {
                let next_chunk = current_chunk.next();
                let prev_chunk = current_chunk.prev();

                debug_assert!(current_chunk.is_empty());
                debug_assert!(next_chunk.is_empty());
                debug_assert!(next_chunk.next().is_dead());

                if !next_chunk.is_dead() {
                    let (smaller, bigger) = if current_chunk.size() > next_chunk.size() {
                        (next_chunk, current_chunk)
                    } else {
                        (current_chunk, next_chunk)
                    };
                    current_chunk = bigger;
                    current_chunk.set_prev(prev_chunk);
                    current_chunk.set_next(ChunkRef::dead());
                    self.drop_chunk(smaller);
                }

                let inlined_chunk = self.inlined.get_mut();
                if prev_chunk.is_dead() {
                    inlined_chunk.set_next_chunk(current_chunk);
                    if const { N > 0 } {
                        inlined_chunk.set_current_chunk(ChunkRef::dead());
                        inlined_chunk.get().dealloc_element()
                    } else {
                        inlined_chunk.set_current_chunk(current_chunk);
                        None
                    }
                } else {
                    prev_chunk.set_next(current_chunk);
                    inlined_chunk.set_current_chunk(prev_chunk);
                    prev_chunk.dealloc_element()
                }
            }
        }
    }

    /// Allocates a new chunk of memory with at least the given rough size, and
    /// updates the capacity of the vector accordingly.
    unsafe fn new_chunk(&self, rough_size: usize) -> ChunkRef<T> {
        let (new_chunk, capacity) = ChunkRef::new(rough_size);
        self.cap_or_len.update(|cap| cap + capacity);
        new_chunk
    }

    /// Drops all elements in the chunk and deallocates it.
    unsafe fn drop_chunk(&mut self, chunk: ChunkRef<T>) {
        debug_assert!(!T::IS_ZST);
        let chunk_capacity = unsafe { chunk.capacity() };
        unsafe { chunk.drop() };
        self.cap_or_len.update(|cap| cap - chunk_capacity);
    }
}

#[cfg(test)]
mod utest;