cfr/
split.rs

1//! Split slices into uneven chunks
2use std::iter::FusedIterator;
3use std::mem;
4
5pub struct SplitsBy<'a, T, I> {
6    slice: &'a [T],
7    lens: I,
8}
9
10impl<'a, T, I: Iterator<Item = usize>> Iterator for SplitsBy<'a, T, I> {
11    type Item = &'a [T];
12
13    fn next(&mut self) -> Option<Self::Item> {
14        match self.lens.next() {
15            Some(len) => {
16                let (ret, rest) = self.slice.split_at(len);
17                self.slice = rest;
18                Some(ret)
19            }
20            None => None,
21        }
22    }
23}
24
25impl<T, I: FusedIterator<Item = usize>> FusedIterator for SplitsBy<'_, T, I> {}
26
27pub fn split_by<T, I: Iterator<Item = usize>, N: IntoIterator<IntoIter = I>>(
28    slice: &[T],
29    lens: N,
30) -> SplitsBy<'_, T, I> {
31    SplitsBy {
32        slice,
33        lens: lens.into_iter(),
34    }
35}
36
37pub struct SplitsByMut<'a, T, I> {
38    slice: &'a mut [T],
39    lens: I,
40}
41
42impl<'a, T, I: Iterator<Item = usize>> Iterator for SplitsByMut<'a, T, I> {
43    type Item = &'a mut [T];
44
45    fn next(&mut self) -> Option<Self::Item> {
46        match self.lens.next() {
47            Some(len) => {
48                let tmp = mem::take(&mut self.slice);
49                let (ret, rest) = tmp.split_at_mut(len);
50                self.slice = rest;
51                Some(ret)
52            }
53            None => None,
54        }
55    }
56}
57
58impl<T, I: FusedIterator<Item = usize>> FusedIterator for SplitsByMut<'_, T, I> {}
59
60pub fn split_by_mut<T, I: Iterator<Item = usize>, N: IntoIterator<IntoIter = I>>(
61    slice: &mut [T],
62    lens: N,
63) -> SplitsByMut<'_, T, I> {
64    SplitsByMut {
65        slice,
66        lens: lens.into_iter(),
67    }
68}
69
70#[cfg(test)]
71mod tests {
72
73    #[test]
74    fn test_ref() {
75        let slices = [1, 2, 3, 4, 5, 6, 7, 8];
76        let sliced: Box<[&[_]]> = super::split_by(&slices, [1, 3, 4]).collect();
77        assert_eq!(*sliced, [vec![1], vec![2, 3, 4], vec![5, 6, 7, 8]]);
78    }
79
80    #[test]
81    fn test_mut() {
82        let mut slices = [1, 2, 3, 4, 5, 6, 7, 8];
83        for mut_slice in super::split_by_mut(&mut slices, [1, 3, 4]) {
84            mut_slice.fill(0);
85        }
86        assert_eq!(slices, [0; 8]);
87    }
88}