pushgen/structs/
take.rs

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

pushgen/structs/
take.rs
