chunked_vec 0.5.1

use std::iter::FusedIterator;

use crate::{Chunk, ChunkedVec};

/// A mutable iterator over the logical chunks of a [`ChunkedVec`].
///
/// This struct is created by the [`iter_chunks_mut`] method on [`ChunkedVec`].
pub struct ChunksMut<'a, T, const N: usize> {
    pub(crate) inner: std::slice::IterMut<'a, Chunk<T, N>>,
    pub(crate) remaining_len: usize,
}

impl<T, const N: usize> ChunkedVec<T, N> {
    /// Returns an iterator over the logical chunks of the vector as mutable
    /// slices.
    ///
    /// Each item contains the initialized elements of one chunk. The last
    /// chunk may be shorter than `N`.
    ///
    /// # Examples
    /// ```
    /// use chunked_vec::{ChunkedVec, ChunkedVecSized};
    ///
    /// let mut vec: ChunkedVec<i32, 4> = ChunkedVecSized::new();
    /// vec.extend(0..6);
    ///
    /// for chunk in vec.iter_chunks_mut() {
    ///     for value in chunk {
    ///         *value *= 2;
    ///     }
    /// }
    ///
    /// assert_eq!(vec.chunk(0), Some(&[0, 2, 4, 6][..]));
    /// assert_eq!(vec.chunk(1), Some(&[8, 10][..]));
    /// ```
    pub fn iter_chunks_mut(&mut self) -> ChunksMut<'_, T, N> {
        let chunk_count = self.chunk_count();
        let len = self.len();

        ChunksMut {
            inner: self.data[..chunk_count].iter_mut(),
            remaining_len: len,
        }
    }
}

impl<'a, T, const N: usize> Iterator for ChunksMut<'a, T, N> {
    type Item = &'a mut [T];

    fn next(&mut self) -> Option<Self::Item> {
        let chunk = self.inner.next()?;
        let len = self.remaining_len.min(N);
        self.remaining_len -= len;

        unsafe {
            Some(std::slice::from_raw_parts_mut(
                chunk.as_mut_ptr().cast::<T>(),
                len,
            ))
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.inner.len();
        (remaining, Some(remaining))
    }
}

impl<T, const N: usize> ExactSizeIterator for ChunksMut<'_, T, N> {}

impl<T, const N: usize> FusedIterator for ChunksMut<'_, T, N> {}

#[cfg(test)]
mod tests {
    use crate::{ChunkedVec, ChunkedVecSized};

    #[test]
    fn test_iter_chunks_mut() {
        let mut vec: ChunkedVec<i32, 4> = ChunkedVecSized::new();
        vec.extend(0..10);

        for chunk in vec.iter_chunks_mut() {
            for value in chunk {
                *value += 10;
            }
        }

        let chunks: Vec<&[i32]> = vec.iter_chunks().collect();
        assert_eq!(
            chunks,
            vec![&[10, 11, 12, 13][..], &[14, 15, 16, 17][..], &[18, 19][..]]
        );
    }

    #[test]
    fn test_iter_chunks_mut_empty() {
        let mut vec: ChunkedVec<i32, 4> = ChunkedVecSized::new();
        assert_eq!(vec.iter_chunks_mut().next(), None);
    }
}