pushgen/structs/
skip.rs

1use crate::{Generator, GeneratorResult, ValueResult};
2use core::num::NonZeroUsize;
3
4/// Skip over a set amount of values. See [`.skip()`](crate::GeneratorExt::skip) for more details.
5#[derive(Clone)]
6pub struct Skip<Gen> {
7    generator: Gen,
8    amount: usize,
9}
10
11impl<Gen> Skip<Gen> {
12    #[inline]
13    pub(crate) fn new(generator: Gen, amount: usize) -> Self {
14        Self { generator, amount }
15    }
16}
17
18impl<Gen> Generator for Skip<Gen>
19where
20    Gen: Generator,
21{
22    type Output = Gen::Output;
23
24    #[inline]
25    fn run(&mut self, mut output: impl FnMut(Self::Output) -> ValueResult) -> GeneratorResult {
26        if self.amount > 0 {
27            // Safety: checked by if clause
28            match self
29                .generator
30                .try_advance(unsafe { core::num::NonZeroUsize::new_unchecked(self.amount) })
31            {
32                (_, GeneratorResult::Complete) => {
33                    self.amount = 0;
34                    return GeneratorResult::Complete;
35                }
36                (x, _) => {
37                    self.amount -= x;
38                    if self.amount != 0 {
39                        return GeneratorResult::Stopped;
40                    }
41                }
42            }
43        }
44
45        self.generator.run(|value| output(value))
46    }
47
48    #[inline]
49    fn try_advance(&mut self, n: NonZeroUsize) -> (usize, GeneratorResult) {
50        if self.amount > 0 {
51            match self
52                .generator
53                .try_advance(unsafe { NonZeroUsize::new_unchecked(self.amount) })
54            {
55                (_, GeneratorResult::Complete) => {
56                    self.amount = 0;
57                    return (0, GeneratorResult::Complete);
58                }
59                (x, _) => {
60                    self.amount -= x;
61                    if self.amount != 0 {
62                        return (0, GeneratorResult::Stopped);
63                    }
64                }
65            }
66        }
67
68        self.generator.try_advance(n)
69    }
70}
71
72#[derive(Clone)]
73pub struct SkipWhile<Src, P> {
74    source: Src,
75    predicate: P,
76    need_skip_run: bool,
77}
78
79impl<Src, P> SkipWhile<Src, P>
80where
81    Src: Generator,
82    P: FnMut(&Src::Output) -> bool,
83{
84    #[inline]
85    pub(crate) fn new(source: Src, predicate: P) -> Self {
86        Self {
87            source,
88            predicate,
89            need_skip_run: true,
90        }
91    }
92}
93
94impl<Src, P> Generator for SkipWhile<Src, P>
95where
96    Src: Generator,
97    P: FnMut(&Src::Output) -> bool,
98{
99    type Output = Src::Output;
100
101    #[inline]
102    fn run(&mut self, mut output: impl FnMut(Self::Output) -> ValueResult) -> GeneratorResult {
103        if self.need_skip_run {
104            let mut first_to_push = None;
105            let predicate = &mut self.predicate;
106
107            let skip_run_result = self.source.run(|x| {
108                if predicate(&x) {
109                    ValueResult::MoreValues
110                } else {
111                    first_to_push = Some(x);
112                    ValueResult::Stop
113                }
114            });
115
116            if skip_run_result == GeneratorResult::Complete {
117                return GeneratorResult::Complete;
118            } else if let Some(x) = first_to_push {
119                self.need_skip_run = false;
120                if output(x) == ValueResult::Stop {
121                    return GeneratorResult::Stopped;
122                }
123            } else {
124                return GeneratorResult::Stopped;
125            }
126        }
127        self.source.run(|x| output(x))
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134    use crate::test::StoppingGen;
135    use crate::{GeneratorExt, IntoGenerator};
136
137    #[test]
138    fn skip() {
139        let a = [1, 2, 3];
140
141        let mut gen = Skip::new(a.into_gen(), 2);
142        let mut output = Vec::new();
143        let result = gen.for_each(|x| output.push(x));
144        assert_eq!(output, [&3]);
145        assert_eq!(result, GeneratorResult::Complete);
146    }
147
148    #[test]
149    fn spuriously_stopping_skip() {
150        let data = [1, 2, 3, 4, 5];
151        for x in 0..5 {
152            let mut gen = StoppingGen::new(x, &data).skip(3);
153            let mut output = Vec::new();
154            let result = gen.for_each(|x| output.push(x));
155            assert_eq!(result, GeneratorResult::Stopped);
156            let result = gen.for_each(|x| output.push(x));
157            assert_eq!(result, GeneratorResult::Complete);
158            assert_eq!(output, [&4, &5]);
159        }
160    }
161
162    #[test]
163    fn skip_while() {
164        let a = [-1i32, 0, 1];
165
166        let mut gen = SkipWhile::new(a.into_gen(), |x| x.is_negative());
167        let mut output = Vec::new();
168        let result = gen.for_each(|x| output.push(x));
169        assert_eq!(output, [&0, &1]);
170        assert_eq!(result, GeneratorResult::Complete);
171    }
172
173    #[test]
174    fn spuriously_stopping_skip_while() {
175        let data = [-1i32, -2, 0, -1, 2];
176        for x in 0..5 {
177            let mut gen = StoppingGen::new(x, &data).skip_while(|x| x.is_negative());
178            let mut output = Vec::new();
179            let result = gen.for_each(|x| output.push(x));
180            assert_eq!(result, GeneratorResult::Stopped);
181            let result = gen.for_each(|x| output.push(x));
182            assert_eq!(result, GeneratorResult::Complete);
183            assert_eq!(output, [&0, &-1, &2]);
184        }
185    }
186
187    #[test]
188    fn try_advance() {
189        let data = [0, 1, 2, 3, 4, 5, 6, 7, 8];
190        let mut gen = data.into_gen().skip(3);
191        let result = gen.try_advance(NonZeroUsize::new(2).unwrap());
192        assert_eq!(result, (2, GeneratorResult::Stopped));
193        assert_eq!(gen.next(), Ok(&5));
194    }
195
196    #[test]
197    fn try_advance_stopping_skip_region() {
198        let data = [0, 1, 2, 3, 4, 5, 6, 7, 8];
199        let mut gen = StoppingGen::new(1, &data).skip(3);
200        let result = gen.try_advance(NonZeroUsize::new(2).unwrap());
201        assert_eq!(result, (0, GeneratorResult::Stopped));
202        let result = gen.try_advance(NonZeroUsize::new(2).unwrap());
203        assert_eq!(result, (2, GeneratorResult::Stopped));
204        assert_eq!(gen.next(), Ok(&5));
205    }
206
207    #[test]
208    fn try_advance_stopping() {
209        let data = [0, 1, 2, 3, 4, 5, 6, 7, 8];
210        let mut gen = StoppingGen::new(4, &data).skip(3);
211        let result = gen.try_advance(NonZeroUsize::new(2).unwrap());
212        assert_eq!(result, (1, GeneratorResult::Stopped));
213        let result = gen.try_advance(NonZeroUsize::new(1).unwrap());
214        assert_eq!(result, (1, GeneratorResult::Stopped));
215        assert_eq!(gen.next(), Ok(&5));
216    }
217
218    #[test]
219    fn try_advance_max() {
220        let data = [0, 1, 2, 3, 4, 5, 6, 7, 8];
221        let mut gen = data.into_gen().skip(3);
222        let result = gen.try_advance(NonZeroUsize::new(usize::MAX).unwrap());
223        assert_eq!(result, (6, GeneratorResult::Complete));
224    }
225}
pushgen/structs/skip.rs

pushgen/structs/
skip.rs
