librashader/librashader-runtime/src/array_chunks_mut.rs

86 lines
2.8 KiB
Rust

/// An iterator over a slice in (non-overlapping) mutable chunks (`N` elements
/// at a time), starting at the beginning of the slice.
///
/// When the slice len is not evenly divided by the chunk size, the last
/// up to `N-1` elements will be omitted but can be retrieved from
/// the [`into_remainder`] function from the iterator.
///
/// This struct is created by the [`array_chunks_mut`] method on [slices].
///
///
/// [`array_chunks_mut`]: slice::array_chunks_mut
/// [`into_remainder`]: ../../std/slice/struct.ArrayChunksMut.html#method.into_remainder
/// [slices]: slice
#[derive(Debug)]
#[must_use = "iterators are lazy and do nothing unless consumed"]
pub struct ArrayChunksMut<'a, T: 'a, const N: usize> {
iter: core::slice::IterMut<'a, [T; N]>,
}
impl<'a, T, const N: usize> ArrayChunksMut<'a, T, N> {
#[inline]
pub(super) fn new(slice: &'a mut [T]) -> Self {
let (array_slice, _rem) = as_chunks_mut(slice);
Self { iter: array_slice.iter_mut() }
}
}
impl<'a, T, const N: usize> Iterator for ArrayChunksMut<'a, T, N> {
type Item = &'a mut [T; N];
#[inline]
fn next(&mut self) -> Option<&'a mut [T; N]> {
self.iter.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
#[inline]
fn count(self) -> usize {
self.iter.count()
}
#[inline]
fn nth(&mut self, n: usize) -> Option<Self::Item> {
self.iter.nth(n)
}
#[inline]
fn last(self) -> Option<Self::Item> {
self.iter.last()
}
}
/// Splits the slice into a slice of `N`-element arrays,
/// starting at the beginning of the slice,
/// and a remainder slice with length strictly less than `N`.
///
/// # Panics
///
/// Panics if `N` is 0. This check will most probably get changed to a compile time
/// error before this method gets stabilized.
///
#[inline]
#[must_use]
fn as_chunks_mut<T, const N: usize>(slice: &mut [T]) -> (&mut [[T; N]], &mut [T]) {
unsafe fn as_chunks_unchecked_mut<T, const N: usize>(slice: &mut [T]) -> &mut [[T; N]] {
// SAFETY: Caller must guarantee that `N` is nonzero and exactly divides the slice length
let new_len = slice.len() / N;
// SAFETY: We cast a slice of `new_len * N` elements into
// a slice of `new_len` many `N` elements chunks.
unsafe { core::slice::from_raw_parts_mut(slice.as_mut_ptr().cast(), new_len) }
}
assert!(N != 0, "chunk size must be non-zero");
let len = slice.len() / N;
let (multiple_of_n, remainder) = slice.split_at_mut(len * N);
// SAFETY: We already panicked for zero, and ensured by construction
// that the length of the subslice is a multiple of N.
let array_slice = unsafe { as_chunks_unchecked_mut(multiple_of_n) };
(array_slice, remainder)
}