mod.rs

// #![allow(unused)]

use std::ops::{Generator, GeneratorState};
use std::pin::Pin;

// mod nightly_gen;
// Define a wrapper struct for the generator trait
pub struct Gen<'a, G, T>
where
    G: FnMut() -> Box<dyn Generator<Yield = T, Return = ()> + 'a + Unpin>, // Generator mutable function signature
    T: 'static, // Lifetime bound (might be able to use something shorter)
{
    generator_fn: G, // The generator function
    generator: Option<Pin<Box<dyn Generator<Yield = T, Return = ()> + 'a>>>, // The generator instance
}

impl<'a, G, T> Gen<'a, G, T>
where
    G: FnMut() -> Box<dyn Generator<Yield = T, Return = ()> + 'a + Unpin>,
    T: 'static,
{
    // Constructor for the generator wrapper
    pub fn new(generator_fn: G) -> Self {
        Gen {
            generator_fn,
            generator: None, // Initialize the generator instance to None
        }
    }
}

// Implementation of the Iterator trait for the generator wrapper
impl<'a, G, T> Iterator for Gen<'a, G, T>
where
    G: FnMut() -> Box<dyn Generator<Yield = T, Return = ()> + 'a + Unpin>,
    T: 'static,
{
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        // If the generator instance is None, create it using the provided generator function
        if self.generator.is_none() {
            self.generator = Some(Box::pin((self.generator_fn)()));
        }

        // Extract mutable reference (can't be None because of line above, so should never throw a panic)
        let generator = self.generator.as_mut().unwrap();

        // Resume the generator and match its state
        match generator.as_mut().resume(()) {
            GeneratorState::Yielded(value) => Some(value), // Yielded a value
            GeneratorState::Complete(_) => None,           // Generator is complete
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn gen_inside_gen() {
        let vecs_gen = Gen::new(|| {
            Box::new(move || {
                let fib_gen = Gen::new(|| {
                    let n: i32 = 200;
                    Box::new(move || {
                        let mut i = 1;
                        let mut j = 1;
                        while i < n {
                            yield i;
                            i = i + j;
                            j = i - j;
                        }
                    })
                });

                for i in fib_gen {
                    let mut v = Vec::new();
                    for j in i..i + 10 {
                        v.push(j);
                    }
                    yield v;
                }
            })
        });

        let mut result: Vec<Vec<i32>> = Vec::new();
        for i in vecs_gen {
            result.push(i);
        }

        assert_eq!(
            result,
            vec![
                vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
                vec![2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
                vec![3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
                vec![5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
                vec![8, 9, 10, 11, 12, 13, 14, 15, 16, 17],
                vec![13, 14, 15, 16, 17, 18, 19, 20, 21, 22],
                vec![21, 22, 23, 24, 25, 26, 27, 28, 29, 30],
                vec![34, 35, 36, 37, 38, 39, 40, 41, 42, 43],
                vec![55, 56, 57, 58, 59, 60, 61, 62, 63, 64],
                vec![89, 90, 91, 92, 93, 94, 95, 96, 97, 98],
                vec![144, 145, 146, 147, 148, 149, 150, 151, 152, 153],
            ]
        );
    }
    #[test]
    fn factorial() {
        let fact_gen = Gen::new(|| {
            Box::new(move || {
                let mut n = 1;
                let mut f = 1;
                while n <= 10 {
                    yield f;
                    n = n + 1;
                    f = f * n;
                }
            })
        });

        let mut result = Vec::new();
        for f in fact_gen {
            result.push(f);
        }

        assert_eq!(
            result,
            vec![1, 2, 6, 24, 120, 720, 5040, 40320, 362880, 3628800]
        );
    }

    #[test]
    fn no_yield() {
        let empty = Gen::new(|| {
            Box::new(|| {
                return; // TODO what to do here? Keep it? Is it nonsense?
                yield ();
            })
        });

        let mut result = Vec::new();
        for i in empty {
            result.push(i);
        }
        assert!(result == vec![]);
    }

    #[test]
    fn only_one() {
        let single_gen = Gen::new(|| {
            Box::new(|| {
                yield i32::MIN; // yield only one value
            })
        });

        let mut result = Vec::new();
        for i in single_gen {
            result.push(i);
        }

        assert_eq!(result, vec![i32::MIN]);
    }

    #[test]
    fn infinite() {
        let infinite_gen = Gen::new(|| {
            Box::new(|| {
                let mut n = 1;
                loop {
                    yield n; // yield a value in each iteration
                    n = n + 1;
                }
            })
        });

        let mut result = Vec::new();
        for i in infinite_gen {
            if i > 10 {
                break;
            }
            result.push(i);
        }

        assert_eq!(result, vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
    }

    #[test]
    fn gen_box() {
        let boxes_gen = Gen::new(|| {
            Box::new(|| {
                yield Box::new("Hello");
                yield Box::new("Bye");
            })
        });

        let mut result = Vec::new();
        for i in boxes_gen {
            result.push(*i);
        }

        assert_eq!(result, vec!["Hello", "Bye"]);
    }

    use std::cell::RefCell;
    use std::rc::Rc;

    #[test]
    fn gen_rcs() {
        let rcs_gen = Gen::new(|| {
            Box::new(move || {
                yield Rc::new(RefCell::new([1, 2, 3]));
                yield Rc::new(RefCell::new([4, 5, 6]));
                yield Rc::new(RefCell::new([7, 8, 9]));
            })
        });

        let mut result = Vec::new();
        for rc in rcs_gen {
            result.push((*rc).take());
        }

        assert_eq!(result, vec![[1, 2, 3], [4, 5, 6], [7, 8, 9]]);
    }

    #[test]
    fn gen_empty_tuple() {
        let void_tuple_gen = Gen::new(|| {
            Box::new(|| {
                yield ();
                yield ();
            })
        });

        let mut result = Vec::new();
        for _ in void_tuple_gen {
            result.push(());
        }

        assert_eq!(result, vec![(), ()]);
    }

    #[test]
    fn closure() {
        fn fibonacci(n: u64) -> u64 {
            let fib = Gen::new(|| {
                Box::new(|| {
                    let mut i = 1;
                    let mut j = 1;
                    while i < n {
                        yield i;
                        i = i + j;
                        j = i - j;
                    }
                })
            });
            let mut result = 0;
            for i in fib {
                result += i;
            }
            result
        }
        assert![fibonacci(10_000) == 17709];
    }
}