bump-stack 0.5.0

use crate::util::{self, TypeProps};
use crate::{Chunk, ChunkFooter, Stack, wrap};
use core::iter::{DoubleEndedIterator, Iterator};
use core::marker::PhantomData;
use core::ptr::{self, NonNull};
use core::slice;

macro_rules! impl_iter {
    (struct $iter_name:ident {
        *$ptr_mut:tt,
        &{ $( $ref_mut:tt )? },
        $as_ref:ident,
        $without_provenance:ident,
        $(,)?
    }) => {
        pub struct $iter_name<'a, T> {
            /// The chunk's footer where `first_or_idx` besides within.
            first_footer: NonNull<ChunkFooter>,

            /// The chunk's footer where `last_or_len` besides within.
            last_footer: NonNull<ChunkFooter>,

            /// The address of an element that was pushed first, and should be returned
            /// by this iterator last.
            first: *$ptr_mut T,

            /// # For non-ZST elements
            ///
            /// The address of element that was pushed last, and should be returned by
            /// this iterator first.
            ///
            /// # For ZST elements
            ///
            /// The number of elements that the iterator should run over.
            last_or_len: *$ptr_mut T,

            phantom: PhantomData<&'a $( $ref_mut )? T>,
        }

        impl<'a, T> $iter_name<'a, T> {
            pub(crate) fn new(stack: &'a $( $ref_mut )? Stack<T>) -> Self {
                if const { T::IS_ZST } {
                    Self {
                        first_footer: stack.current_footer.get(),
                        last_footer: stack.current_footer.get(),
                        first: ptr::$without_provenance(0),
                        last_or_len: ptr::$without_provenance(stack.len()),
                        phantom: PhantomData,
                    }
                } else {
                    let first_ptr = unsafe { wrap::<T>(stack.first_footer.get()).end_aligned() };
                    let current_chunk = unsafe { wrap::<T>(stack.current_footer.get()) };
                    let last_ptr = current_chunk.ptr().cast().as_ptr();
                    Self {
                        first_footer: stack.first_footer.get(),
                        last_footer: current_chunk.footer(),
                        first: first_ptr.cast().as_ptr(),
                        last_or_len: last_ptr,
                        phantom: PhantomData,
                    }
                }
            }
        }

        impl<'a, T> $iter_name<'a, T> {
            #[inline(always)]
            unsafe fn next_element_fast(&mut self) -> Option<NonNull<T>> {
                let ptr = self.last_or_len as usize;
                let end = self.last_footer.as_ptr() as usize;
                debug_assert!(ptr <= end);

                let chunk_is_empty = if const { Chunk::<T>::FOOTER_IS_END } {
                    ptr == end
                } else {
                    end - ptr < T::SIZE
                };

                if !chunk_is_empty {
                    let ptr = self.last_or_len;
                    self.last_or_len = ptr.wrapping_add(1);
                    Some(unsafe { NonNull::new_unchecked(ptr as *mut T) })
                } else {
                    None
                }
            }

            unsafe fn next_element_slow(&mut self) -> Option<NonNull<T>> {
                unsafe {
                    let last_chunk = wrap::<T>(self.last_footer);
                    self.last_footer = last_chunk.prev();

                    let new_last_chunk = wrap::<T>(self.last_footer);
                    self.last_or_len = new_last_chunk.ptr().cast().as_ptr();

                    if self.first == self.last_or_len {
                        return None;
                    }

                    self.next_element_fast()
                }
            }

            #[inline(always)]
            unsafe fn prev_element_fast(&mut self) -> Option<NonNull<T>> {
                unsafe {
                    let first_chunk = wrap::<T>(self.first_footer);
                    let chunk_start = first_chunk.start().cast().as_ptr() as *$ptr_mut T;
                    debug_assert!(chunk_start <= self.first);

                    let ptr = self.first;
                    if ptr != chunk_start {
                        self.first = ptr.wrapping_sub(1);
                        debug_assert!(chunk_start <= self.first);
                        Some(NonNull::new_unchecked(self.first as *mut T))
                    } else {
                        None
                    }
                }
            }

            unsafe fn prev_element_slow(&mut self) -> Option<NonNull<T>> {
                unsafe {
                    let first_chunk = wrap::<T>(self.first_footer);
                    self.first_footer = first_chunk.next();
                    let new_first_chunk = wrap::<T>(self.first_footer);
                    self.first = new_first_chunk.end_aligned().cast().as_ptr();

                    if self.first == self.last_or_len {
                        return None;
                    }

                    self.prev_element_fast()
                }
            }
        }

        impl<'a, T> Iterator for $iter_name<'a, T> {
            type Item = &'a $( $ref_mut )? T;

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

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

            #[inline]
            fn count(self) -> usize
            where
                Self: Sized,
            {
                if const { T::IS_ZST } {
                    self.last_or_len as usize
                } else {
                    let mut count = 0;
                    let buffer_end = self.last_footer.as_ptr() as usize;
                    let ptr = self.last_or_len as usize;
                    count += (buffer_end - ptr) / T::SIZE;

                    let mut chunk = unsafe { wrap::<T>(wrap::<T>(self.last_footer).prev()) };
                    while !chunk.is_dead() {
                        count += chunk.capacity();
                        chunk = unsafe { wrap(chunk.prev()) };
                    }
                    count
                }
            }

            fn nth(&mut self, n: usize) -> Option<Self::Item> {
                let m = n + 1;
                let mut need = m * T::SIZE;
                if const { T::IS_ZST } {
                    let length = self.last_or_len as usize;
                    if length > n {
                        unsafe {
                            self.last_or_len = self.last_or_len.wrapping_byte_sub(m);
                            Some(util::zst_ptr().$as_ref())
                        }
                    } else {
                        self.last_or_len = ptr::$without_provenance(0);
                        None
                    }
                } else {
                    loop {
                        let last_chunk = unsafe { wrap::<T>(self.last_footer) };
                        let first_ptr = self.first;
                        let last_ptr = self.last_or_len;

                        if self.last_footer == self.first_footer {
                            debug_assert!(last_ptr <= first_ptr);
                            let size = first_ptr as usize - last_ptr as usize;
                            if need <= size {
                                let new_ptr = last_ptr.wrapping_byte_add(need);
                                self.last_or_len = new_ptr;
                                let ptr = last_ptr.wrapping_byte_add(need - T::SIZE);
                                break Some(unsafe { &$( $ref_mut )? *ptr });
                            } else {
                                self.last_or_len = first_ptr;
                                break None;
                            }
                        } else {
                            // can ignore the gap between the buffer and the footer,
                            // because: 0 <= gap < ELEMENT_SIZE
                            let chunk_end = last_chunk.end().as_ptr();
                            let size = chunk_end as usize - last_ptr as usize;
                            if need <= size {
                                let new_ptr = last_ptr.wrapping_byte_add(need);
                                self.last_or_len = new_ptr;
                                let ptr = last_ptr.wrapping_byte_add(need - T::SIZE);
                                break Some(unsafe { &$( $ref_mut )? *ptr });
                            } else {
                                need -= size;
                                self.last_footer = last_chunk.prev();
                                let last_chunk = unsafe { wrap(self.last_footer) };
                                self.last_or_len = last_chunk.start().as_ptr();
                            }
                        }
                    }
                }
            }

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

        impl<'a, T> DoubleEndedIterator for $iter_name<'a, T> {
            #[inline]
            fn next_back(&mut self) -> Option<Self::Item> {
                if const { T::IS_ZST } {
                    // it just decrements `last_or_len`
                    return self.next();
                }
                if self.first == self.last_or_len {
                    return None;
                }
                unsafe {
                    if let Some($( $ref_mut )? elem_ptr) = self.prev_element_fast() {
                        Some(elem_ptr.$as_ref())
                    } else {
                        self.prev_element_slow().map(|$( $ref_mut )? ptr| ptr.$as_ref())
                    }
                }
            }
        }
    };
}

impl_iter!(
    struct Iter {
        *const,
        &{},
        as_ref,
        without_provenance,
    }
);

impl_iter!(
    struct IterMut {
        *mut,
        &{ mut },
        as_mut,
        without_provenance_mut,
    }
);

macro_rules! impl_slice_iter {
    (struct $iter_name:ident {
        *$ptr_mut:tt,
        &{ $( $ref_mut:tt )? },
        $ptr_null:ident,
        $from_raw_parts:ident,
        $without_provenance:ident,
        $chunk_slice:ident,
        $(,)?
    }) => {
        pub struct $iter_name<'a, T> {
            first_footer: NonNull<ChunkFooter>,
            last_footer: NonNull<ChunkFooter>,
            ptr_or_len: *$ptr_mut T,
            phantom: PhantomData<&'a $( $ref_mut )? [T]>,
        }

        impl<'a, T> $iter_name<'a, T> {
            pub(crate) fn new(stack: &'a $( $ref_mut )? Stack<T>) -> Self {
                let first = stack.first_footer.get();

                let mut last = stack.current_footer.get();
                let mut last_chunk = unsafe { wrap::<T>(last) };

                let last_ptr = if const { T::IS_ZST } {
                    ptr::$without_provenance(stack.len())
                } else {
                    // asdf
                    if stack.is_empty() {
                        debug_assert_eq!(first, last);
                        ptr::$ptr_null()
                    } else {
                        if last_chunk.is_empty() {
                            last = last_chunk.prev();
                            last_chunk = unsafe { wrap::<T>(last) };
                        }
                        last_chunk.ptr().as_ptr()
                    }
                };
                Self {
                    first_footer: first,
                    last_footer: last,
                    ptr_or_len: last_ptr,
                    phantom: PhantomData,
                }
            }

            fn next_crossed(&mut self) -> Option<&'a $( $ref_mut )? [T]> {
                debug_assert!(!T::IS_ZST);
                debug_assert_eq!(self.first_footer, self.last_footer);

                if !self.ptr_or_len.is_null() {
                    let chunk = unsafe { wrap::<T>(self.last_footer) };

                    debug_assert!(chunk.start().as_ptr().cast_const() <= self.ptr_or_len);
                    debug_assert!(self.ptr_or_len <= chunk.end().cast().as_ptr());

                    let len = (chunk.end().as_ptr() as usize - self.ptr_or_len as usize) / T::SIZE;
                    let slice = unsafe { slice::$from_raw_parts(self.ptr_or_len, len) };
                    self.ptr_or_len = ptr::$ptr_null();
                    Some(slice)
                } else {
                    None
                }
            }

            fn next_zst(&mut self) -> Option<&'a $( $ref_mut )? [T]> {
                debug_assert!(T::IS_ZST);

                let len = self.ptr_or_len as usize;
                if len > 0 {
                    let ptr = util::zst_ptr().as_ptr();
                    self.ptr_or_len = ptr::$without_provenance(0);
                    Some(unsafe { slice::$from_raw_parts(ptr, len) })
                } else {
                    None
                }
            }
        }

        impl<'a, T> Iterator for $iter_name<'a, T> {
            type Item = &'a $( $ref_mut )? [T];

            fn next(&mut self) -> Option<Self::Item> {
                if const { T::IS_ZST } {
                    return self.next_zst();
                }

                if self.first_footer != self.last_footer {
                    let chunk = unsafe { wrap::<T>(self.first_footer) };
                    self.first_footer = chunk.next();
                    Some(unsafe { chunk.$chunk_slice() })
                } else {
                    self.next_crossed()
                }
            }
        }

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

                if self.first_footer != self.last_footer {
                    let chunk = unsafe { wrap::<T>(self.last_footer) };

                    debug_assert!(chunk.start().as_ptr().cast_const() <= self.ptr_or_len);
                    debug_assert!(self.ptr_or_len <= chunk.end().cast().as_ptr());

                    let len = (chunk.end().as_ptr() as usize - self.ptr_or_len as usize) / T::SIZE;
                    let slice = unsafe { slice::$from_raw_parts(self.ptr_or_len, len) };
                    self.last_footer = chunk.prev();
                    self.ptr_or_len = unsafe { wrap(self.last_footer).ptr().as_ptr() };
                    Some(slice)
                } else {
                    self.next_crossed()
                }
            }
        }
    };
}

impl_slice_iter!(
    struct ChunkIter {
        *const,
        &{},
        null,
        from_raw_parts,
        without_provenance,
        slice,
    }
);

impl_slice_iter!(
    struct ChunkIterMut {
        *mut,
        &{ mut },
        null_mut,
        from_raw_parts_mut,
        without_provenance_mut,
        slice_mut,
    }
);

pub struct IntoIter<T> {
    stack: Stack<T>,
}

impl<T> IntoIter<T> {
    #[inline]
    pub(crate) fn new(stack: Stack<T>) -> Self {
        Self { stack }
    }
}

impl<T> Iterator for IntoIter<T> {
    type Item = T;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.stack.pop()
    }
}