computation_process/
generator.rs

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