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