wutil/iter/
split.rs

1/// Iterator returned by [`Iterator::split`](crate::prelude::IterCloneExt::split)
2pub struct Split<I, F>
3where
4    I: Iterator + Clone,
5    F: FnMut(&I::Item) -> bool + Clone,
6{
7    iter: Option<I>,
8    pred: F,
9}
10
11impl<I, F> Split<I, F>
12where
13    I: Iterator + Clone,
14    F: FnMut(&I::Item) -> bool + Clone,
15{
16    pub(super) fn new(iter: I, pred: F) -> Self {
17        Self {
18            iter: Some(iter),
19            pred,
20        }
21    }
22}
23
24/// Iterator used by [`Iterator::split`](crate::prelude::IterCloneExt::split)
25pub struct SplitIterator<I, F>
26where
27    I: Iterator + Clone,
28    F: FnMut(&I::Item) -> bool + Clone,
29{
30    iter: Option<I>,
31    pred: F,
32}
33
34impl<I, F> Iterator for SplitIterator<I, F>
35where
36    I: Iterator + Clone,
37    F: FnMut(&I::Item) -> bool + Clone,
38{
39    type Item = I::Item;
40
41    fn next(&mut self) -> Option<Self::Item> {
42        let item = self.iter.as_mut()?.next()?;
43
44        if (self.pred)(&item) {
45            self.iter = None;
46            return None;
47        }
48
49        Some(item)
50    }
51}
52
53impl<I, F> Iterator for Split<I, F>
54where
55    I: Iterator + Clone,
56    F: FnMut(&I::Item) -> bool + Clone,
57{
58    type Item = SplitIterator<I, F>;
59
60    fn next(&mut self) -> Option<Self::Item> {
61        let next_iter = self.iter.clone()?;
62
63        loop {
64            let Some(item) = self.iter.as_mut()?.next() else {
65                self.iter = None;
66                break;
67            };
68
69            if (self.pred)(&item) {
70                break;
71            }
72        }
73
74        Some(SplitIterator {
75            iter: Some(next_iter),
76            pred: self.pred.clone(),
77        })
78    }
79}
80
81/// Iterator returned by [`Iterator::split_inclusive`](crate::prelude::IterCloneExt::split_inclusive)
82pub struct SplitInclusive<I, F>
83where
84    I: Iterator + Clone,
85    F: FnMut(&I::Item) -> bool + Clone,
86{
87    iter: Option<I>,
88    pred: F,
89}
90
91impl<I, F> SplitInclusive<I, F>
92where
93    I: Iterator + Clone,
94    F: FnMut(&I::Item) -> bool + Clone,
95{
96    pub(super) fn new(iter: I, pred: F) -> Self {
97        Self {
98            iter: Some(iter),
99            pred,
100        }
101    }
102}
103
104/// Iterator used by [`Iterator::split_inclusive`](crate::prelude::IterCloneExt::split_inclusive)
105pub struct InclusiveSplitIterator<I, F>
106where
107    I: Iterator + Clone,
108    F: FnMut(&I::Item) -> bool + Clone,
109{
110    iter: Option<I>,
111    pred: F,
112}
113
114impl<I, F> Iterator for InclusiveSplitIterator<I, F>
115where
116    I: Iterator + Clone,
117    F: FnMut(&I::Item) -> bool + Clone,
118{
119    type Item = I::Item;
120
121    fn next(&mut self) -> Option<Self::Item> {
122        let item = self.iter.as_mut()?.next()?;
123
124        if (self.pred)(&item) {
125            self.iter = None;
126        }
127
128        Some(item)
129    }
130}
131
132impl<I, F> Iterator for SplitInclusive<I, F>
133where
134    I: Iterator + Clone,
135    F: FnMut(&I::Item) -> bool + Clone,
136{
137    type Item = InclusiveSplitIterator<I, F>;
138
139    fn next(&mut self) -> Option<Self::Item> {
140        let next_iter = self.iter.clone()?;
141
142        loop {
143            let Some(item) = self.iter.as_mut()?.next() else {
144                self.iter = None;
145                break;
146            };
147
148            if (self.pred)(&item) {
149                break;
150            }
151        }
152
153        Some(InclusiveSplitIterator {
154            iter: Some(next_iter),
155            pred: self.pred.clone(),
156        })
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use crate::prelude::*;
163    use alloc::vec::Vec;
164
165    #[test]
166    fn split() {
167        let tests: &[(&[u32], &[&[u32]])] = &[
168            (
169                &[0, 10, 20, 0, 0, 50, 0],
170                &[&[], &[10, 20], &[], &[50], &[]],
171            ),
172            (
173                &[0, 10, 20, 0, 0, 50, 0, 2],
174                &[&[], &[10, 20], &[], &[50], &[2]],
175            ),
176        ];
177
178        for (input, expected) in tests {
179            let split_nums: Vec<Vec<u32>> = input
180                .iter()
181                .split(|n| **n == 0)
182                .map(|n| n.copied().collect::<Vec<u32>>())
183                .collect();
184
185            assert_eq!(&split_nums, expected);
186        }
187    }
188
189    #[test]
190    fn split_inclusive() {
191        let tests: &[(&[u32], &[&[u32]])] = &[
192            (
193                &[0, 10, 20, 0, 0, 50, 0],
194                &[&[0], &[10, 20, 0], &[0], &[50, 0], &[]],
195            ),
196            (
197                &[0, 10, 20, 0, 0, 50, 0, 2],
198                &[&[0], &[10, 20, 0], &[0], &[50, 0], &[2]],
199            ),
200        ];
201
202        for (input, expected_output) in tests {
203            let split_nums: Vec<Vec<u32>> = input
204                .iter()
205                .split_inclusive(|n| **n == 0)
206                .map(|n| n.copied().collect::<Vec<u32>>())
207                .collect();
208
209            assert_eq!(&split_nums, expected_output);
210        }
211    }
212}