use core::iter::{FusedIterator, Iterator};
use core::marker::Unpin;
use core::ops::{Coroutine, CoroutineState};
use core::pin::Pin;
#[derive(Copy, Clone, Debug)]
pub struct GenIter<T>(pub T)
where
T: Coroutine<Return = ()> + Unpin;
impl<T> Iterator for GenIter<T>
where
T: Coroutine<Return = ()> + Unpin,
{
type Item = T::Yield;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
match Pin::new(&mut self.0).resume(()) {
CoroutineState::Yielded(n) => Some(n),
CoroutineState::Complete(()) => None,
}
}
}
impl<G> From<G> for GenIter<G>
where
G: Coroutine<Return = ()> + Unpin,
{
#[inline]
fn from(gen: G) -> Self {
GenIter(gen)
}
}
#[macro_export]
macro_rules! gen_iter {
($block: block) => {
$crate::GenIter(
#[coroutine]
|| $block,
)
};
(move $block: block) => {
$crate::GenIter(
#[coroutine]
move || $block,
)
};
}
#[derive(Copy, Clone, Debug)]
pub struct GenIterReturn<G: Coroutine + Unpin>(Result<G::Return, G>);
impl<G: Coroutine + Unpin> GenIterReturn<G> {
#[inline]
pub fn new(g: G) -> Self {
GenIterReturn(Err(g))
}
#[inline]
pub fn is_done(&self) -> bool {
self.0.is_ok()
}
#[inline]
pub fn return_or_self(self) -> Result<G::Return, Self> {
match self.0 {
Ok(r) => Ok(r),
Err(_) => Err(self),
}
}
}
impl<G: Coroutine + Unpin> Iterator for &mut GenIterReturn<G> {
type Item = G::Yield;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
match self.0 {
Ok(_) => None,
Err(ref mut g) => match Pin::new(g).resume(()) {
CoroutineState::Yielded(y) => Some(y),
CoroutineState::Complete(r) => {
self.0 = Ok(r);
None
}
},
}
}
}
impl<G: Coroutine + Unpin> FusedIterator for &mut GenIterReturn<G> {}
impl<G: Coroutine + Unpin> From<G> for GenIterReturn<G> {
#[inline]
fn from(g: G) -> Self {
GenIterReturn::new(g)
}
}
#[macro_export]
macro_rules! gen_iter_return {
($block: block) => {
$crate::GenIterReturn::new(
#[coroutine]
|| $block,
)
};
(move $block: block) => {
$crate::GenIterReturn::new(
#[coroutine]
move || $block,
)
};
}
#[cfg(test)]
mod tests {
use super::GenIterReturn;
#[test]
fn gen_iter_works() {
let mut g = gen_iter!({
yield 1;
yield 2;
});
assert_eq!(g.next(), Some(1));
assert_eq!(g.next(), Some(2));
assert_eq!(g.next(), None);
}
#[test]
fn gen_iter_macro() {
let mut g = gen_iter!(move {
yield 1;
yield 2;
});
assert_eq!(g.next(), Some(1));
assert_eq!(g.next(), Some(2));
assert_eq!(g.next(), None);
}
#[test]
fn gen_iter_return_works() {
let mut g = GenIterReturn::new(
#[coroutine]
|| {
yield 1;
"done"
},
);
assert_eq!((&mut g).next(), Some(1));
assert!(!g.is_done());
g = match g.return_or_self() {
Ok(_) => panic!("coroutine is done but should not"),
Err(g) => g,
};
assert_eq!((&mut g).next(), None);
assert!(g.is_done());
assert_eq!((&mut g).next(), None);
assert_eq!(g.return_or_self().ok(), Some("done"));
}
#[test]
fn from_coroutine() {
let mut g = GenIterReturn::from(
#[coroutine]
|| {
yield 1;
"done"
},
);
assert_eq!((&mut g).next(), Some(1));
assert_eq!((&mut g).next(), None);
assert!(g.is_done());
assert_eq!(g.return_or_self().ok(), Some("done"));
}
#[test]
fn macro_usage() {
let mut g = gen_iter_return!(move {
yield 1;
yield 2;
return "done";
});
let (mut sum, mut count) = (0, 0);
for y in &mut g {
sum += y;
count += 1;
}
assert_eq!((sum, count), (3, 2));
assert!(g.is_done());
assert_eq!(g.return_or_self().ok(), Some("done"));
}
}