std-next 0.1.9

Opinionated utilities and polyfills
Documentation
#![allow(unsafe_code)]

use crate::sealed::Sealed;

pub trait SliceExt<T>: Sealed {
    fn get2_mut(&mut self, i: usize, j: usize) -> Option<(&mut T, &mut T)>;

    fn as_chunks_<const N: usize>(&self) -> (&[[T; N]], &[T]);

    fn as_chunks_mut_<const N: usize>(&mut self) -> (&mut [[T; N]], &mut [T]);
}

impl<T> Sealed for [T] {}

impl<T> SliceExt<T> for [T] {
    fn get2_mut(&mut self, i: usize, j: usize) -> Option<(&mut T, &mut T)> {
        if i == j || i >= self.len() || j >= self.len() {
            return None;
        }
        // SAFETY: `i` and `j` are in bounds and not equal
        unsafe {
            let ptr = self.as_mut_ptr();
            Some((&mut *ptr.add(i), &mut *ptr.add(j)))
        }
    }

    fn as_chunks_<const N: usize>(&self) -> (&[[T; N]], &[T]) {
        assert!(N > 0, "chunk size must be non-zero");
        let base = self.as_ptr();
        let len = self.len();
        let (div, rem) = (len / N, len % N);
        let chunks = unsafe { slice(base.cast(), div) };
        let rest = unsafe { slice(base.add(div * N), rem) };
        (chunks, rest)
    }

    fn as_chunks_mut_<const N: usize>(&mut self) -> (&mut [[T; N]], &mut [T]) {
        assert!(N > 0, "chunk size must be non-zero");
        let base = self.as_mut_ptr();
        let len = self.len();
        let (div, rem) = (len / N, len % N);
        let chunks = unsafe { slice_mut(base.cast(), div) };
        let rest = unsafe { slice_mut(base.add(div * N), rem) };
        (chunks, rest)
    }
}

unsafe fn slice<'a, T>(data: *const T, len: usize) -> &'a [T] {
    core::slice::from_raw_parts(data, len)
}

unsafe fn slice_mut<'a, T>(data: *mut T, len: usize) -> &'a mut [T] {
    core::slice::from_raw_parts_mut(data, len)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_as_chunks() {
        let mut data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
        let (chunks, rest) = data.as_chunks_::<3>();
        assert_eq!(chunks, [[1, 2, 3], [4, 5, 6], [7, 8, 9]]);
        assert_eq!(rest, [10]);

        let (chunks, rest) = data.as_chunks_::<4>();
        assert_eq!(chunks, [[1, 2, 3, 4], [5, 6, 7, 8]]);
        assert_eq!(rest, [9, 10]);

        let (chunks, rest) = data.as_chunks_::<10>();
        assert_eq!(chunks, [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]);
        assert_eq!(rest, &[]);

        let (chunks, rest) = data.as_chunks_mut_::<3>();
        assert_eq!(chunks, [[1, 2, 3], [4, 5, 6], [7, 8, 9]]);
        assert_eq!(rest, [10]);

        let (chunks, rest) = data.as_chunks_mut_::<4>();
        assert_eq!(chunks, [[1, 2, 3, 4], [5, 6, 7, 8]]);
        assert_eq!(rest, [9, 10]);

        let (chunks, rest) = data.as_chunks_mut_::<10>();
        assert_eq!(chunks, [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]);
        assert_eq!(rest, &[]);
    }
}