use core::ops::GeneratorState;
use core::pin::Pin;
use super::*;
#[derive(PartialEq, Eq)]
struct TestItemA {}
#[derive(PartialEq, Eq)]
struct TestItemB {}
struct TestIterator {
count: usize,
}
impl TestIterator {
pub fn new() -> Self {
Self { count: 0 }
}
}
impl Iterator for TestIterator {
type Item = TestItemA;
fn next(&mut self) -> Option<Self::Item> {
if self.count < 2 {
self.count += 1;
Some(TestItemA {})
} else {
None
}
}
}
struct TestGenerator {
count: usize,
}
impl TestGenerator {
pub fn new() -> Self {
Self { count: 0 }
}
}
impl Generator<()> for TestGenerator {
type Yield = TestItemA;
type Return = TestItemB;
fn resume(self: Pin<&mut Self>, _arg: ()) -> GeneratorState<Self::Yield, Self::Return> {
let me = self.get_mut();
if me.count < 2 {
me.count += 1;
GeneratorState::Yielded(TestItemA {})
} else {
GeneratorState::Complete(TestItemB {})
}
}
}
#[test]
fn iter_as_gen() {
let mut gen = &mut TestIterator::new().into_generator();
assert!(gen.resume(()) == GeneratorState::Yielded(TestItemA {}));
assert!(gen.resume(()) == GeneratorState::Yielded(TestItemA {}));
assert!(gen.resume(()) == GeneratorState::Complete(()));
}
#[test]
fn gen_as_states_iter() {
let mut gen = TestGenerator::new();
let mut states = gen.states();
let mut iter = states.iter();
assert!(iter.next() == Some(GeneratorState::Yielded(TestItemA {})));
assert!(iter.next() == Some(GeneratorState::Yielded(TestItemA {})));
assert!(iter.next() == Some(GeneratorState::Complete(TestItemB {})));
assert!(iter.next() == None);
}