computation_process/
generator.rs

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
7/// Defines a single step of a [`Generator`].
8///
9/// Implement this trait to define the logic for generating items.
10/// Each call to `step` should either:
11/// - Return `Ok(Some(item))` to yield an item
12/// - Return `Ok(None)` when the generator is exhausted
13/// - Return `Err(Incomplete::Suspended)` to yield control without producing an item
14/// - Return `Err(Incomplete::Cancelled(_))` if cancellation was detected
15///
16/// # Type Parameters
17///
18/// - `CONTEXT`: Immutable configuration/input for the generator
19/// - `STATE`: Mutable state that persists across steps
20/// - `ITEM`: The type of items produced by the generator
21pub trait GeneratorStep<CONTEXT, STATE, ITEM> {
22    /// Execute one step of the generator.
23    ///
24    /// Returns `Some(item)` to yield an item, or `None` when exhausted.
25    fn step(context: &CONTEXT, state: &mut STATE) -> Completable<Option<ITEM>>;
26}
27
28/// A stateful generator that can be suspended and resumed.
29///
30/// `Generator` is the default implementation of [`GenAlgorithm`]. It delegates the
31/// actual generation logic to a [`GeneratorStep`] implementation while handling
32/// the boilerplate of state management and cancellation checking.
33///
34/// `Generator` implements both [`Generatable`] (for suspendable iteration) and
35/// [`Iterator`] (for convenient collection, skipping suspensions automatically).
36///
37/// # Type Parameters
38///
39/// - `CONTEXT`: Immutable configuration passed to each step
40/// - `STATE`: Mutable state that persists across steps
41/// - `ITEM`: The type of items produced
42/// - `STEP`: The [`GeneratorStep`] implementation that defines the generation logic
43///
44/// # Example
45///
46/// ```rust
47/// use computation_process::{Generator, GeneratorStep, Completable, Generatable, Stateful};
48///
49/// struct CountStep;
50///
51/// impl GeneratorStep<u32, u32, u32> for CountStep {
52///     fn step(max: &u32, current: &mut u32) -> Completable<Option<u32>> {
53///         *current += 1;
54///         if *current <= *max {
55///             Ok(Some(*current))
56///         } else {
57///             Ok(None)
58///         }
59///     }
60/// }
61///
62/// let mut generator = Generator::<u32, u32, u32, CountStep>::from_parts(3, 0);
63/// assert_eq!(generator.try_next(), Some(Ok(1)));
64/// assert_eq!(generator.try_next(), Some(Ok(2)));
65/// assert_eq!(generator.try_next(), Some(Ok(3)));
66/// assert_eq!(generator.try_next(), None);
67/// ```
68#[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        // After 3 items, should return None
233        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        // The first call should suspend (state becomes 1)
285        assert_eq!(generator.try_next(), Some(Err(Incomplete::Suspended)));
286        assert_eq!(*generator.state(), 1);
287
288        // The second call should suspend (state becomes 2)
289        assert_eq!(generator.try_next(), Some(Err(Incomplete::Suspended)));
290        assert_eq!(*generator.state(), 2);
291
292        // Third call should return item (state becomes 3)
293        let item = generator.try_next().unwrap().unwrap();
294        assert_eq!(item, 3);
295        assert_eq!(*generator.state(), 3);
296
297        // Fourth call should return item (state becomes 4)
298        let item = generator.try_next().unwrap().unwrap();
299        assert_eq!(item, 4);
300
301        // The fifth call should return None
302        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        // Iterator should skip suspensions automatically
310        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        // This will generate infinite items since state never changes
346        // But we can test that it generates at least one
347        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            // The first call says "exhausted"; subsequent calls would yield an item.
356            // A correct Generator must not call the step again after it observes exhaustion.
357            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}