bumpish 0.4.0

use super::RawBump;
use super::chunk::ChunkRef;
use crate::util::{self, TypeProps};
use core::marker::PhantomData;
use core::num::NonZero;
use core::ptr::{self, NonNull};

pub(crate) struct IterPtr<'a, T, const N: usize> {
    first_chunk: ChunkRef<T>,
    last_chunk: ChunkRef<T>,
    start: *mut T,
    end_or_len: *mut T,
    phantom: PhantomData<&'a RawBump<T, N>>,
}

impl<'a, T, const N: usize> IterPtr<'a, T, N> {
    /// Creates an iterator over the elements of the bump allocator.
    pub(crate) fn new(bump: &'a RawBump<T, N>) -> Self {
        if const { T::IS_ZST } {
            Self {
                first_chunk: ChunkRef::dead(),
                last_chunk: ChunkRef::dead(),
                start: ptr::without_provenance_mut(0),
                end_or_len: ptr::without_provenance_mut(unsafe { bump.zst_len() }),
                phantom: PhantomData,
            }
        } else if const { N == 0 } {
            let inlined_chunk = unsafe { bump.inlined_ptr().as_ref() };
            let first_chunk = inlined_chunk.next_chunk();
            let last_chunk = inlined_chunk.current_chunk();
            Self {
                first_chunk,
                last_chunk,
                start: unsafe { first_chunk.start().as_ptr() },
                end_or_len: unsafe { last_chunk.ptr().as_ptr() },
                phantom: PhantomData,
            }
        } else {
            unsafe {
                let mut inlined_chunk_ptr = bump.inlined_ptr();
                let first_chunk = ChunkRef::from(inlined_chunk_ptr.cast());
                let inlined_chunk = inlined_chunk_ptr.as_mut();
                let current_chunk = inlined_chunk.current_chunk();
                let last_chunk = if current_chunk.is_dead() {
                    first_chunk
                } else {
                    current_chunk
                };
                Self {
                    first_chunk,
                    last_chunk,
                    start: first_chunk.start().as_ptr(),
                    end_or_len: last_chunk.ptr().as_ptr(),
                    phantom: PhantomData,
                }
            }
        }
    }

    #[inline(always)]
    unsafe fn next_element_fast(&mut self) -> Option<NonNull<T>> {
        unsafe {
            if self.start != self.first_chunk.end().as_ptr() {
                let old_start = self.start;
                self.start = old_start.add(1);
                Some(NonNull::new_unchecked(old_start))
            } else {
                None
            }
        }
    }

    unsafe fn next_chunk(&mut self) {
        debug_assert!(!self.first_chunk.is(self.last_chunk));
        unsafe {
            self.first_chunk = self.first_chunk.next();
            self.start = self.first_chunk.start().as_ptr();
        }
    }

    unsafe fn next_element_slow(&mut self) -> Option<NonNull<T>> {
        unsafe {
            self.next_chunk();
            if self.start != self.end_or_len {
                Some(self.next_element_fast().unwrap_unchecked())
            } else {
                None
            }
        }
    }

    #[inline(always)]
    unsafe fn prev_element_fast(&mut self) -> Option<NonNull<T>> {
        unsafe {
            if self.end_or_len != self.last_chunk.start().as_ptr() {
                self.end_or_len = self.end_or_len.sub(1);
                Some(NonNull::new_unchecked(self.end_or_len))
            } else {
                None
            }
        }
    }

    unsafe fn prev_chunk(&mut self) {
        debug_assert!(!self.first_chunk.is(self.last_chunk));
        unsafe {
            self.last_chunk = self.last_chunk.prev();
            if self.last_chunk.is_dead() {
                // it means that the first chunk is inlined chunk
                self.last_chunk = self.first_chunk;
            }
            self.end_or_len = self.last_chunk.ptr().as_ptr();
        }
    }

    unsafe fn prev_element_slow(&mut self) -> Option<NonNull<T>> {
        unsafe {
            self.prev_chunk();
            if self.start != self.end_or_len {
                Some(self.prev_element_fast().unwrap_unchecked())
            } else {
                None
            }
        }
    }

    fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
        if const { T::IS_ZST } {
            let len = self.end_or_len as usize;
            if let Some(k) = len.checked_sub(n) {
                self.end_or_len = ptr::without_provenance_mut(k);
                Ok(())
            } else {
                self.end_or_len = ptr::without_provenance_mut(0);
                Err(unsafe { NonZero::new_unchecked(n - len) })
            }
        } else {
            let mut n_size = if let Some(n_size) = n.checked_mul(T::SIZE) {
                n_size
            } else {
                return Err(unsafe { NonZero::new_unchecked(n - self.count()) });
            };
            loop {
                let start = self.start;
                let end = self.end_or_len;

                if self.first_chunk.is(self.last_chunk) {
                    debug_assert!(start <= end);
                    let data_size = end as usize - start as usize;
                    if n_size <= data_size {
                        self.start = unsafe { start.byte_add(n_size) };
                        break Ok(());
                    } else {
                        self.start = end;
                        let k = (n_size - data_size) / T::SIZE;
                        break Err(unsafe { NonZero::new_unchecked(k) });
                    }
                } else {
                    let data_size =
                        unsafe { self.first_chunk.end().as_ptr() } as usize - start as usize;
                    if n_size <= data_size {
                        self.start = unsafe { start.byte_add(n_size) };
                        break Ok(());
                    } else {
                        n_size -= data_size;
                        unsafe { self.next_chunk() };
                    }
                }
            }
        }
    }

    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
        if const { T::IS_ZST } {
            self.advance_by(n)
        } else {
            let mut n_size = if let Some(n_size) = n.checked_mul(T::SIZE) {
                n_size
            } else {
                return Err(unsafe { NonZero::new_unchecked(n - self.count()) });
            };
            loop {
                let start = self.start;
                let end = self.end_or_len;

                if self.first_chunk.is(self.last_chunk) {
                    debug_assert!(start <= end);
                    let data_size = end as usize - start as usize;
                    if n_size <= data_size {
                        self.end_or_len = unsafe { end.byte_sub(n_size) };
                        break Ok(());
                    } else {
                        self.end_or_len = start;
                        let k = (n_size - data_size) / T::SIZE;
                        break Err(unsafe { NonZero::new_unchecked(k) });
                    }
                } else {
                    let data_size =
                        end as usize - unsafe { self.last_chunk.start().as_ptr() } as usize;
                    if n_size <= data_size {
                        self.end_or_len = unsafe { end.byte_sub(n_size) };
                        break Ok(());
                    } else {
                        n_size -= data_size;
                        unsafe { self.prev_chunk() };
                    }
                }
            }
        }
    }

    fn count(&self) -> usize
    where
        Self: Sized,
    {
        if const { T::IS_ZST } {
            return self.end_or_len as usize;
        }

        if self.first_chunk.is(self.last_chunk) {
            debug_assert!(self.start <= self.end_or_len);
            return (self.end_or_len as usize - self.start as usize) / T::SIZE;
        }

        let first_ptr = self.start as usize;
        let first_end = unsafe { self.first_chunk.end().as_ptr() as usize };
        let mut size = first_end - first_ptr;

        unsafe {
            let mut chunk = self.first_chunk.next();
            while !chunk.is(self.last_chunk) {
                size += chunk.data_size();
                chunk = chunk.next();
            }
        }
        let last_ptr = self.end_or_len as usize;
        let last_start = unsafe { self.last_chunk.start().as_ptr() as usize };
        size += last_ptr - last_start;
        size / T::SIZE
    }
}

impl<'a, T, const N: usize> Iterator for IterPtr<'a, T, N> {
    type Item = NonNull<T>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if const { T::IS_ZST } {
            if self.end_or_len as usize != 0 {
                self.end_or_len = self.end_or_len.wrapping_byte_sub(1);
                Some(util::zst_ptr())
            } else {
                None
            }
        } else {
            if self.start == self.end_or_len {
                return None;
            }
            unsafe {
                if let Some(elem_ptr) = self.next_element_fast() {
                    Some(elem_ptr)
                } else {
                    self.next_element_slow()
                }
            }
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        if const { T::IS_ZST } {
            let remains = self.end_or_len as usize;
            (remains, Some(remains))
        } else {
            (0, None)
        }
    }

    fn count(self) -> usize
    where
        Self: Sized,
    {
        Self::count(&self)
    }

    fn last(mut self) -> Option<Self::Item> {
        self.next_back()
    }

    fn nth(&mut self, n: usize) -> Option<Self::Item> {
        self.advance_by(n).ok()?;
        self.next()
    }
}

impl<'a, T, const N: usize> DoubleEndedIterator for IterPtr<'a, T, N> {
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        if const { T::IS_ZST } {
            self.next()
        } else {
            if self.start == self.end_or_len {
                return None;
            }
            unsafe {
                if let Some(elem_ptr) = self.prev_element_fast() {
                    Some(elem_ptr)
                } else {
                    self.prev_element_slow()
                }
            }
        }
    }

    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
        self.advance_back_by(n).ok()?;
        self.next_back()
    }
}

pub(crate) struct ChunkIterPtr<'a, T, const N: usize> {
    first_chunk: ChunkRef<T>,
    last_chunk: ChunkRef<T>,
    ptr_or_len: *mut T,
    phantom: PhantomData<&'a RawBump<T, N>>,
}

impl<'a, T, const N: usize> ChunkIterPtr<'a, T, N> {
    /// Creates an iterator over the chunks of the bump allocator.
    pub(super) fn new(core: &'a RawBump<T, N>) -> Self {
        if const { T::IS_ZST } {
            let len = unsafe { core.zst_len() };
            Self {
                first_chunk: ChunkRef::dead(),
                last_chunk: ChunkRef::dead(),
                ptr_or_len: ptr::without_provenance_mut(len),
                phantom: PhantomData,
            }
        } else {
            let first_chunk = if const { N > 0 } {
                unsafe { ChunkRef::from(core.inlined_ptr().cast()) }
            } else {
                unsafe { core.inlined_ptr().as_ref() }.next_chunk()
            };
            let mut current_chunk = unsafe { core.inlined_ptr().as_ref() }.current_chunk();
            if unsafe { current_chunk.is_empty() } {
                current_chunk = unsafe { current_chunk.prev() };
            }
            let last_chunk = if unsafe { current_chunk.is_dead() } {
                first_chunk
            } else {
                current_chunk
            };
            let ptr_or_len = if unsafe { first_chunk.is_empty() } {
                ptr::null_mut()
            } else {
                unsafe { last_chunk.ptr().as_ptr() }
            };
            Self {
                first_chunk,
                last_chunk,
                ptr_or_len,
                phantom: PhantomData,
            }
        }
    }

    unsafe fn next_zst(&mut self) -> Option<(NonNull<T>, usize)> {
        debug_assert!(T::IS_ZST);
        let len = self.ptr_or_len as usize;
        if len != 0 {
            self.ptr_or_len = ptr::without_provenance_mut(0);
            Some((util::zst_ptr(), len))
        } else {
            None
        }
    }

    fn next_crossed(&mut self) -> Option<(NonNull<T>, usize)> {
        debug_assert!(!T::IS_ZST);
        debug_assert!(self.first_chunk.is(self.last_chunk));
        if !self.ptr_or_len.is_null() {
            let start = unsafe { self.first_chunk.start() };
            let len = (self.ptr_or_len as usize - start.as_ptr() as usize) / T::SIZE;
            self.ptr_or_len = ptr::null_mut();
            Some((start, len))
        } else {
            None
        }
    }
}

impl<'a, T, const N: usize> Iterator for ChunkIterPtr<'a, T, N> {
    type Item = (NonNull<T>, usize);

    fn next(&mut self) -> Option<Self::Item> {
        if const { T::IS_ZST } {
            unsafe { self.next_zst() }
        } else {
            if !self.first_chunk.is(self.last_chunk) {
                let slice = unsafe { self.first_chunk.slice_ptr() };
                self.first_chunk = unsafe { self.first_chunk.next() };
                Some(slice)
            } else {
                self.next_crossed()
            }
        }
    }
}

impl<'a, T, const N: usize> DoubleEndedIterator for ChunkIterPtr<'a, T, N> {
    fn next_back(&mut self) -> Option<Self::Item> {
        if const { T::IS_ZST } {
            return unsafe { self.next_zst() };
        }

        if !self.first_chunk.is(self.last_chunk) {
            let start = unsafe { self.last_chunk.start() };
            let end = self.ptr_or_len;
            let len = (end as usize - start.as_ptr() as usize) / T::SIZE;
            self.last_chunk = unsafe { self.last_chunk.prev() };
            if unsafe { self.last_chunk.is_dead() } {
                self.last_chunk = self.first_chunk;
            }
            self.ptr_or_len = unsafe { self.last_chunk.ptr().as_ptr() };
            Some((start, len))
        } else {
            self.next_crossed()
        }
    }
}

pub(crate) struct IntoIter<T, const N: usize> {
    bump: RawBump<T, N>,

    first_chunk: ChunkRef<T>,
    last_chunk: ChunkRef<T>,

    /// offset within `first_chunk`
    start: usize,

    /// offset inside `last_chunk`, and number of elements it `T` is ZST
    end: usize,
}

impl<T, const N: usize> IntoIter<T, N> {
    pub(crate) fn new(mut bump: RawBump<T, N>) -> Self {
        if const { T::IS_ZST } {
            Self {
                bump,
                first_chunk: ChunkRef::dead(),
                last_chunk: ChunkRef::dead(),
                start: 0,
                end: 0,
            }
        } else if const { N == 0 } {
            let inlined_chunk = bump.inlined.get_mut();
            let first_chunk = inlined_chunk.next_chunk();
            let last_chunk = inlined_chunk.current_chunk();
            Self {
                bump,
                first_chunk,
                last_chunk,
                start: unsafe { first_chunk.start_offset() },
                end: unsafe { last_chunk.ptr_offset() },
            }
        } else {
            unsafe {
                let inlined_chunk = bump.inlined.get_mut();
                let first_chunk = ChunkRef::dead();
                let start = inlined_chunk.get().start_offset();
                let current_chunk = inlined_chunk.current_chunk();
                let (last_chunk, end) = if current_chunk.is_dead() {
                    (first_chunk, inlined_chunk.get().ptr_offset())
                } else {
                    (current_chunk, current_chunk.ptr_offset())
                };
                Self {
                    bump,
                    first_chunk,
                    last_chunk,
                    start,
                    end,
                }
            }
        }
    }

    unsafe fn first_chunk(&mut self) -> ChunkRef<T> {
        debug_assert!(!T::IS_ZST);
        unsafe {
            if const { N > 0 } && self.first_chunk.is_dead() {
                self.bump.inlined.get_mut().get()
            } else {
                self.first_chunk
            }
        }
    }

    unsafe fn last_chunk(&mut self) -> ChunkRef<T> {
        debug_assert!(!T::IS_ZST);
        unsafe {
            if const { N > 0 } && self.last_chunk.is_dead() {
                self.bump.inlined.get_mut().get()
            } else {
                self.last_chunk
            }
        }
    }

    #[inline(always)]
    unsafe fn next_element_fast(&mut self) -> Option<NonNull<T>> {
        unsafe {
            let chunk = self.first_chunk();
            let end = if self.first_chunk.is(self.last_chunk) {
                self.end
            } else {
                chunk.ptr_offset()
            };

            debug_assert!(end <= chunk.end_offset());
            debug_assert!(self.start <= end);

            if self.start != end {
                let ptr = chunk.ptr_by_offset(self.start);
                self.start += T::SIZE;
                Some(ptr)
            } else {
                None
            }
        }
    }

    unsafe fn next_element_slow(&mut self) -> Option<NonNull<T>> {
        unsafe {
            if !self.first_chunk.is(self.last_chunk) {
                self.next_chunk();
                self.next_element_fast()
            } else {
                None
            }
        }
    }

    unsafe fn next_chunk(&mut self) {
        unsafe {
            debug_assert!(!self.first_chunk.is(self.last_chunk));

            let fake_chunk = self.first_chunk;
            let real_chunk = self.first_chunk();
            self.first_chunk = real_chunk.next();
            self.start = self.first_chunk.start_offset();

            if const { N == 0 } || !fake_chunk.is_dead() {
                // i.e. the prev_chunk is not the inlined one
                self.dealloc_chunk(real_chunk);
            } else {
                real_chunk.reset_ptr();
            }
        }
    }

    #[inline(always)]
    unsafe fn prev_element_fast(&mut self) -> Option<NonNull<T>> {
        unsafe {
            let chunk = self.last_chunk();
            let start = if self.first_chunk.is(self.last_chunk) {
                self.start
            } else {
                chunk.start_offset()
            };

            debug_assert!(chunk.start_offset() <= start);
            debug_assert!(start <= self.end);

            if self.end != start {
                self.end -= T::SIZE;
                Some(chunk.ptr_by_offset(self.end))
            } else {
                None
            }
        }
    }

    unsafe fn prev_element_slow(&mut self) -> Option<NonNull<T>> {
        unsafe {
            if !self.first_chunk.is(self.last_chunk) {
                self.prev_chunk();
                self.prev_element_fast()
            } else {
                None
            }
        }
    }

    unsafe fn prev_chunk(&mut self) {
        unsafe {
            debug_assert!(!self.first_chunk.is(self.last_chunk));
            debug_assert!(!self.last_chunk.is_dead());

            let prev_chunk = self.last_chunk;
            self.last_chunk = self.last_chunk.prev();
            self.end = self.last_chunk().ptr_offset();
            self.dealloc_chunk(prev_chunk);
        }
    }

    /// Deallocates a chunk, updating the bump's capacity as needed.
    unsafe fn dealloc_chunk(&mut self, chunk: ChunkRef<T>) {
        debug_assert!(!T::IS_ZST);
        unsafe {
            let chunk_capacity = chunk.capacity();
            chunk.reset_ptr();
            chunk.drop();
            self.bump.cap_or_len.update(|cap| cap - chunk_capacity);
        };
    }
}

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

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if const { T::IS_ZST } {
            self.bump.deallocate().map(|ptr| unsafe { ptr.read() })
        } else {
            unsafe {
                if let Some(ptr) = self.next_element_fast() {
                    Some(ptr.read())
                } else {
                    self.next_element_slow().map(|ptr| ptr.read())
                }
            }
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        if const { T::IS_ZST } {
            let remains = unsafe { self.bump.zst_len() };
            (remains, Some(remains))
        } else {
            (0, None)
        }
    }
}

impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        if const { T::IS_ZST } {
            self.next()
        } else {
            unsafe {
                if let Some(elem_ptr) = self.prev_element_fast() {
                    Some(elem_ptr.read())
                } else {
                    self.prev_element_slow().map(|ptr| ptr.read())
                }
            }
        }
    }
}

impl<T, const N: usize> Drop for IntoIter<T, N> {
    fn drop(&mut self) {
        unsafe {
            if const { !T::IS_ZST } {
                if const { N > 0 } && self.first_chunk.is_dead() {
                    let inlined_chunk = self.bump.inlined.get_mut().get();
                    let start = inlined_chunk.ptr_by_offset(self.start).as_ptr();
                    let end = if self.last_chunk.is_dead() {
                        inlined_chunk.ptr_by_offset(self.end).as_ptr()
                    } else {
                        inlined_chunk.ptr().as_ptr()
                    };
                    debug_assert!(start <= end);

                    let len = (end as usize - start as usize) / T::SIZE;
                    let slice = ptr::slice_from_raw_parts_mut(start, len);
                    ptr::drop_in_place(slice);
                    self.first_chunk = inlined_chunk.next();
                    self.start = self.first_chunk.start_offset();

                    inlined_chunk.reset_ptr();
                }

                if !self.last_chunk.is_dead() {
                    debug_assert!(!self.first_chunk.is_dead());

                    while !self.first_chunk.is(self.last_chunk) {
                        let curr_chunk = self.first_chunk;
                        self.first_chunk = curr_chunk.next();
                        self.start = self.first_chunk.start_offset();
                        self.bump.drop_chunk(curr_chunk);
                    }

                    if self.start < self.end {
                        let len = (self.end - self.start) / T::SIZE;
                        let start = self.first_chunk.ptr_by_offset(self.start);
                        // TODO: replace with `slice::from_mut_ptr_range` when stabilized
                        let slice = ptr::slice_from_raw_parts_mut(start.as_ptr(), len);
                        ptr::drop_in_place(slice);
                    }
                    self.dealloc_chunk(self.first_chunk);
                }

                let inlined_chunk = self.bump.inlined.get_mut().get();
                inlined_chunk.set_next(ChunkRef::dead());
                inlined_chunk.set_prev(ChunkRef::dead());
            }
        }
    }
}

#[cfg(test)]
mod utest;