1use crate::generatable::Generatable;
2use crate::{Completable, GenAlgorithm, Incomplete, Stateful};
3use cancel_this::{Cancellable, is_cancelled};
4use serde::{Deserialize, Serialize};
5use std::marker::PhantomData;
6
7pub trait GeneratorStep<CONTEXT, STATE, ITEM> {
22 fn step(context: &CONTEXT, state: &mut STATE) -> Completable<Option<ITEM>>;
26}
27
28#[derive(Debug, Serialize, Deserialize)]
69#[serde(
70 bound = "CONTEXT: Serialize + for<'a> Deserialize<'a>, STATE: Serialize + for<'a> Deserialize<'a>"
71)]
72pub struct Generator<CONTEXT, STATE, ITEM, STEP: GeneratorStep<CONTEXT, STATE, ITEM>> {
73 context: CONTEXT,
74 state: STATE,
75 exhausted: bool,
76 #[serde(skip)]
77 _phantom: PhantomData<(ITEM, STEP)>,
78}
79
80impl<CONTEXT, STATE, ITEM, STEP: GeneratorStep<CONTEXT, STATE, ITEM>> Iterator
81 for Generator<CONTEXT, STATE, ITEM, STEP>
82{
83 type Item = Cancellable<ITEM>;
84
85 fn next(&mut self) -> Option<Self::Item> {
86 if self.exhausted {
87 return None;
88 }
89 loop {
90 if let Err(e) = is_cancelled!() {
91 return Some(Err(e));
92 }
93
94 match STEP::step(&self.context, &mut self.state) {
95 Ok(None) => {
96 self.exhausted = true;
97 return None;
98 }
99 Ok(Some(item)) => return Some(Ok(item)),
100 Err(Incomplete::Suspended) => continue,
101 Err(Incomplete::Cancelled(c)) => return Some(Err(c)),
102 Err(Incomplete::Exhausted) => {
103 self.exhausted = true;
104 return None;
105 }
106 }
107 }
108 }
109}
110
111impl<CONTEXT, STATE, OUTPUT, STEP: GeneratorStep<CONTEXT, STATE, OUTPUT>> Generatable<OUTPUT>
112 for Generator<CONTEXT, STATE, OUTPUT, STEP>
113{
114 fn try_next(&mut self) -> Option<Completable<OUTPUT>> {
115 if self.exhausted {
116 return None;
117 }
118 if let Err(e) = is_cancelled!() {
119 return Some(Err(Incomplete::Cancelled(e)));
120 }
121 match STEP::step(&self.context, &mut self.state) {
122 Ok(None) => {
123 self.exhausted = true;
124 None
125 }
126 Ok(Some(v)) => Some(Ok(v)),
127 Err(Incomplete::Exhausted) => {
128 self.exhausted = true;
129 None
130 }
131 Err(e) => Some(Err(e)),
132 }
133 }
134}
135
136impl<CONTEXT, STATE, ITEM, STEP: GeneratorStep<CONTEXT, STATE, ITEM>> Stateful<CONTEXT, STATE>
137 for Generator<CONTEXT, STATE, ITEM, STEP>
138{
139 fn from_parts(context: CONTEXT, state: STATE) -> Self
140 where
141 Self: Sized + 'static,
142 {
143 Generator {
144 context,
145 state,
146 exhausted: false,
147 _phantom: Default::default(),
148 }
149 }
150
151 fn into_parts(self) -> (CONTEXT, STATE) {
152 (self.context, self.state)
153 }
154
155 fn context(&self) -> &CONTEXT {
156 &self.context
157 }
158
159 fn state(&self) -> &STATE {
160 &self.state
161 }
162
163 fn state_mut(&mut self) -> &mut STATE {
164 &mut self.state
165 }
166}
167
168impl<CONTEXT, STATE, ITEM, STEP: GeneratorStep<CONTEXT, STATE, ITEM>>
169 GenAlgorithm<CONTEXT, STATE, ITEM> for Generator<CONTEXT, STATE, ITEM, STEP>
170{
171}
172
173#[cfg(test)]
174mod tests {
175 use super::*;
176 use crate::{GenAlgorithm, Generatable, Incomplete, Stateful};
177 use cancel_this::Cancellable;
178
179 struct SimpleGeneratorStep;
180
181 impl GeneratorStep<i32, u32, String> for SimpleGeneratorStep {
182 fn step(context: &i32, state: &mut u32) -> Completable<Option<String>> {
183 *state += 1;
184 if *state <= 3 {
185 Ok(Some(format!("item-{}-{}", context, state)))
186 } else {
187 Ok(None)
188 }
189 }
190 }
191
192 type SimpleTestGenerator = Generator<i32, u32, String, SimpleGeneratorStep>;
193
194 #[test]
195 fn test_generator_from_parts() {
196 let generator = SimpleTestGenerator::from_parts(42, 0);
197 assert_eq!(*generator.context(), 42);
198 assert_eq!(*generator.state(), 0);
199 }
200
201 #[test]
202 fn test_generator_into_parts() {
203 let generator = SimpleTestGenerator::from_parts(100, 5);
204 let (context, state) = generator.into_parts();
205 assert_eq!(context, 100);
206 assert_eq!(state, 5);
207 }
208
209 #[test]
210 fn test_generator_state_mut() {
211 let mut generator = SimpleTestGenerator::from_parts(42, 0);
212 *generator.state_mut() = 10;
213 assert_eq!(*generator.state(), 10);
214 }
215
216 #[test]
217 fn test_generator_try_next() {
218 let mut generator = SimpleTestGenerator::from_parts(42, 0);
219
220 let item1 = generator.try_next().unwrap().unwrap();
221 assert_eq!(item1, "item-42-1");
222 assert_eq!(*generator.state(), 1);
223
224 let item2 = generator.try_next().unwrap().unwrap();
225 assert_eq!(item2, "item-42-2");
226 assert_eq!(*generator.state(), 2);
227
228 let item3 = generator.try_next().unwrap().unwrap();
229 assert_eq!(item3, "item-42-3");
230 assert_eq!(*generator.state(), 3);
231
232 assert_eq!(generator.try_next(), None);
234 }
235
236 #[test]
237 fn test_generator_iterator() {
238 let generator = SimpleTestGenerator::from_parts(42, 0);
239
240 let items: Vec<Cancellable<String>> = generator.collect();
241 assert_eq!(items.len(), 3);
242 assert_eq!(items[0], Ok("item-42-1".to_string()));
243 assert_eq!(items[1], Ok("item-42-2".to_string()));
244 assert_eq!(items[2], Ok("item-42-3".to_string()));
245 }
246
247 #[test]
248 fn test_generator_dyn_generatable() {
249 let generator = SimpleTestGenerator::from_parts(42, 0);
250 let mut dyn_gen = generator.dyn_generatable();
251 let item = dyn_gen.try_next().unwrap().unwrap();
252 assert_eq!(item, "item-42-1");
253 }
254
255 #[test]
256 fn test_generator_dyn_algorithm() {
257 let generator = SimpleTestGenerator::from_parts(42, 0);
258 let mut dyn_algorithm = generator.dyn_algorithm();
259 let item = dyn_algorithm.try_next().unwrap().unwrap();
260 assert_eq!(item, "item-42-1");
261 }
262
263 struct SuspendingGeneratorStep;
264
265 impl GeneratorStep<(), u32, i32> for SuspendingGeneratorStep {
266 fn step(_context: &(), state: &mut u32) -> Completable<Option<i32>> {
267 *state += 1;
268 if *state <= 2 {
269 Err(Incomplete::Suspended)
270 } else if *state <= 4 {
271 Ok(Some(*state as i32))
272 } else {
273 Ok(None)
274 }
275 }
276 }
277
278 type SuspendingTestGenerator = Generator<(), u32, i32, SuspendingGeneratorStep>;
279
280 #[test]
281 fn test_generator_with_suspensions() {
282 let mut generator = SuspendingTestGenerator::from_parts((), 0);
283
284 assert_eq!(generator.try_next(), Some(Err(Incomplete::Suspended)));
286 assert_eq!(*generator.state(), 1);
287
288 assert_eq!(generator.try_next(), Some(Err(Incomplete::Suspended)));
290 assert_eq!(*generator.state(), 2);
291
292 let item = generator.try_next().unwrap().unwrap();
294 assert_eq!(item, 3);
295 assert_eq!(*generator.state(), 3);
296
297 let item = generator.try_next().unwrap().unwrap();
299 assert_eq!(item, 4);
300
301 assert_eq!(generator.try_next(), None);
303 }
304
305 #[test]
306 fn test_generator_iterator_with_suspensions() {
307 let generator = SuspendingTestGenerator::from_parts((), 0);
308
309 let items: Vec<Cancellable<i32>> = generator.collect();
311 assert_eq!(items.len(), 2);
312 assert_eq!(items[0], Ok(3));
313 assert_eq!(items[1], Ok(4));
314 }
315
316 struct EmptyGeneratorStep;
317
318 impl GeneratorStep<(), (), i32> for EmptyGeneratorStep {
319 fn step(_context: &(), _state: &mut ()) -> Completable<Option<i32>> {
320 Ok(None)
321 }
322 }
323
324 #[test]
325 fn test_empty_generator() {
326 let mut generator = Generator::<(), (), i32, EmptyGeneratorStep>::from_parts((), ());
327 assert_eq!(generator.try_next(), None);
328
329 let items: Vec<Cancellable<i32>> = generator.collect();
330 assert_eq!(items.len(), 0);
331 }
332
333 struct SingleItemGeneratorStep;
334
335 impl GeneratorStep<(), (), i32> for SingleItemGeneratorStep {
336 fn step(_context: &(), _state: &mut ()) -> Completable<Option<i32>> {
337 Ok(Some(42))
338 }
339 }
340
341 #[test]
342 fn test_single_item_generator() {
343 let mut generator = Generator::<(), (), i32, SingleItemGeneratorStep>::from_parts((), ());
344
345 let item = generator.try_next().unwrap().unwrap();
348 assert_eq!(item, 42);
349 }
350
351 struct FlakyExhaustionStep;
352
353 impl GeneratorStep<(), bool, i32> for FlakyExhaustionStep {
354 fn step(_context: &(), state: &mut bool) -> Completable<Option<i32>> {
355 if !*state {
358 *state = true;
359 Ok(None)
360 } else {
361 Ok(Some(123))
362 }
363 }
364 }
365
366 #[test]
367 fn test_generator_is_sticky_exhausted_for_try_next() {
368 let mut generator = Generator::<(), bool, i32, FlakyExhaustionStep>::from_parts((), false);
369 assert_eq!(generator.try_next(), None);
370 assert_eq!(generator.try_next(), None);
371 }
372
373 #[test]
374 fn test_generator_is_sticky_exhausted_for_iterator_next() {
375 let mut generator = Generator::<(), bool, i32, FlakyExhaustionStep>::from_parts((), false);
376 assert_eq!(generator.next(), None);
377 assert_eq!(generator.next(), None);
378 }
379}