1use 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}