stdx/
slice.rs

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

stdx/
slice.rs
