stdx/
slice.rs

1#![allow(unsafe_code)]
2
3use crate::sealed::Sealed;
4
5pub trait SliceExt<T>: Sealed {
6    fn get2_mut(&mut self, i: usize, j: usize) -> Option<(&mut T, &mut T)>;
7
8    fn as_chunks_<const N: usize>(&self) -> (&[[T; N]], &[T]);
9
10    fn as_chunks_mut_<const N: usize>(&mut self) -> (&mut [[T; N]], &mut [T]);
11}
12
13impl<T> Sealed for [T] {}
14
15impl<T> SliceExt<T> for [T] {
16    fn get2_mut(&mut self, i: usize, j: usize) -> Option<(&mut T, &mut T)> {
17        if i == j || i >= self.len() || j >= self.len() {
18            return None;
19        }
20        // SAFETY: `i` and `j` are in bounds and not equal
21        unsafe {
22            let ptr = self.as_mut_ptr();
23            Some((&mut *ptr.add(i), &mut *ptr.add(j)))
24        }
25    }
26
27    fn as_chunks_<const N: usize>(&self) -> (&[[T; N]], &[T]) {
28        assert!(N > 0, "chunk size must be non-zero");
29        let base = self.as_ptr();
30        let len = self.len();
31        let (div, rem) = (len / N, len % N);
32        let chunks = unsafe { slice(base.cast(), div) };
33        let rest = unsafe { slice(base.add(div * N), rem) };
34        (chunks, rest)
35    }
36
37    fn as_chunks_mut_<const N: usize>(&mut self) -> (&mut [[T; N]], &mut [T]) {
38        assert!(N > 0, "chunk size must be non-zero");
39        let base = self.as_mut_ptr();
40        let len = self.len();
41        let (div, rem) = (len / N, len % N);
42        let chunks = unsafe { slice_mut(base.cast(), div) };
43        let rest = unsafe { slice_mut(base.add(div * N), rem) };
44        (chunks, rest)
45    }
46}
47
48unsafe fn slice<'a, T>(data: *const T, len: usize) -> &'a [T] {
49    core::slice::from_raw_parts(data, len)
50}
51
52unsafe fn slice_mut<'a, T>(data: *mut T, len: usize) -> &'a mut [T] {
53    core::slice::from_raw_parts_mut(data, len)
54}
55
56#[cfg(test)]
57mod tests {
58    use super::*;
59
60    #[test]
61    fn test_as_chunks() {
62        let mut data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
63        let (chunks, rest) = data.as_chunks_::<3>();
64        assert_eq!(chunks, [[1, 2, 3], [4, 5, 6], [7, 8, 9]]);
65        assert_eq!(rest, [10]);
66
67        let (chunks, rest) = data.as_chunks_::<4>();
68        assert_eq!(chunks, [[1, 2, 3, 4], [5, 6, 7, 8]]);
69        assert_eq!(rest, [9, 10]);
70
71        let (chunks, rest) = data.as_chunks_::<10>();
72        assert_eq!(chunks, [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]);
73        assert_eq!(rest, &[]);
74
75        let (chunks, rest) = data.as_chunks_mut_::<3>();
76        assert_eq!(chunks, [[1, 2, 3], [4, 5, 6], [7, 8, 9]]);
77        assert_eq!(rest, [10]);
78
79        let (chunks, rest) = data.as_chunks_mut_::<4>();
80        assert_eq!(chunks, [[1, 2, 3, 4], [5, 6, 7, 8]]);
81        assert_eq!(rest, [9, 10]);
82
83        let (chunks, rest) = data.as_chunks_mut_::<10>();
84        assert_eq!(chunks, [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]);
85        assert_eq!(rest, &[]);
86    }
87}
stdx/slice.rs

stdx/
slice.rs
