Skip to content

Commit 85724c7

Browse files
committed
add struct GenIterReturn and macro gen_iter_return to iterate over a generator with return value
1 parent df22576 commit 85724c7

File tree

2 files changed

+187
-0
lines changed

2 files changed

+187
-0
lines changed

src/gen_iter_return.rs

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
use core::ops::{Generator, GeneratorState};
2+
use core::iter::{Iterator, FusedIterator};
3+
use core::marker::Unpin;
4+
use core::pin::Pin;
5+
6+
/// `GenIterReturn<G>` is a iterator for generator with return value.
7+
///
8+
/// Differences with `GenIter<G>`:
9+
/// 1. able to get return value of a generator
10+
/// 2. safe to call `next()` after generator is done without panic
11+
/// 3. maybe less efficient than `GenIter<G>`
12+
#[derive(Copy, Clone, Debug)]
13+
pub struct GenIterReturn<G: Generator + Unpin>(Result<G::Return, G>);
14+
15+
impl<G: Generator + Unpin> GenIterReturn<G> {
16+
#[inline]
17+
pub fn new(g: G) -> Self {
18+
GenIterReturn(Err(g))
19+
}
20+
21+
#[inline]
22+
pub fn is_done(&self) -> bool {
23+
self.0.is_ok()
24+
}
25+
26+
#[inline]
27+
pub fn return_or_self(self) -> Result<G::Return, Self> {
28+
match self.0 {
29+
Ok(r) => Ok(r),
30+
Err(_) => Err(self),
31+
}
32+
}
33+
}
34+
35+
impl<G: Generator + Unpin> Iterator for &mut GenIterReturn<G> {
36+
type Item = G::Yield;
37+
38+
#[inline]
39+
fn next(&mut self) -> Option<Self::Item> {
40+
match self.0 {
41+
Ok(_) => None,
42+
Err(ref mut g) => match Pin::new(g).resume(()) {
43+
GeneratorState::Yielded(y) => Some(y),
44+
GeneratorState::Complete(r) => {
45+
self.0 = Ok(r);
46+
None
47+
},
48+
}
49+
}
50+
}
51+
}
52+
53+
impl<G: Generator + Unpin> FusedIterator for &mut GenIterReturn<G> {}
54+
55+
impl<G: Generator + Unpin> From<G> for GenIterReturn<G> {
56+
#[inline]
57+
fn from(g: G) -> Self {
58+
GenIterReturn::new(g)
59+
}
60+
}
61+
62+
/// macro to simplify iterator - via - generator with return value construction
63+
///
64+
/// Examples:
65+
/// ```
66+
/// #![feature(generators)]
67+
///
68+
/// use gen_iter::gen_iter_return;
69+
///
70+
/// let mut g = gen_iter_return!({
71+
/// yield 1;
72+
/// yield 2;
73+
/// return "done";
74+
/// });
75+
///
76+
/// assert_eq!((&mut g).collect::<Vec<_>>(), [1, 2]); // use `&mut g` as an iterator
77+
/// assert_eq!(g.is_done(), true); // check whether generator is done
78+
/// assert_eq!((&mut g).next(), None); // safe to call `next()` after done
79+
/// assert_eq!(g.return_or_self().ok(), Some("done")); // get return value of generator
80+
/// ```
81+
/// We should use `&mut g` in `for` statement, to keep `g` be valid after `for`, so we can get the return value.
82+
/// ```compile_fail
83+
/// #![feature(generators)]
84+
///
85+
/// use gen_iter::gen_iter_return;
86+
///
87+
/// let mut g = gen_iter_return!({
88+
/// yield 1;
89+
/// yield 2;
90+
/// return "done";
91+
/// });
92+
/// for v in g {} // compile failed, should use `&mut g`
93+
/// ```
94+
#[macro_export]
95+
macro_rules! gen_iter_return {
96+
($block: block) => {
97+
$crate::GenIterReturn::new(|| $block)
98+
};
99+
(move $block: block) => {
100+
$crate::GenIterReturn::new(move || $block)
101+
}
102+
}
103+
104+
#[cfg(test)]
105+
mod tests {
106+
use super::GenIterReturn;
107+
108+
#[test]
109+
fn it_works() {
110+
let mut g = gen_iter_return!({
111+
yield 1;
112+
yield 2;
113+
return "done";
114+
});
115+
116+
assert_eq!((&mut g).next(), Some(1));
117+
assert_eq!(g.is_done(), false);
118+
assert_eq!((&mut g).next(), Some(2));
119+
assert_eq!(g.is_done(), false);
120+
assert_eq!((&mut g).next(), None);
121+
assert_eq!(g.is_done(), true);
122+
assert_eq!((&mut g).next(), None);
123+
assert_eq!(g.return_or_self().ok(), Some("done"));
124+
}
125+
126+
#[test]
127+
fn gen_iter_return_from() {
128+
let mut g: GenIterReturn<_> = GenIterReturn::from(|| {
129+
yield 1;
130+
yield 2;
131+
return "done";
132+
});
133+
let mut gi = &mut g;
134+
135+
assert_eq!(gi.next(), Some(1));
136+
assert_eq!(gi.next(), Some(2));
137+
assert_eq!(gi.next(), None);
138+
139+
assert_eq!(g.is_done(), true);
140+
assert_eq!(g.return_or_self().ok(), Some("done"));
141+
}
142+
143+
#[test]
144+
fn gen_iter_return_macro() {
145+
let mut g = gen_iter_return!({
146+
yield 1;
147+
yield 2;
148+
return "done";
149+
});
150+
151+
let mut sum = 0;
152+
let mut count = 0;
153+
for y in &mut g {
154+
sum += y;
155+
count += 1;
156+
}
157+
assert_eq!(sum, 3);
158+
assert_eq!(count, 2);
159+
160+
assert_eq!(g.is_done(), true);
161+
assert_eq!(g.return_or_self().ok(), Some("done"));
162+
}
163+
}

src/lib.rs

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,35 @@
2828
//! println!("{}", elem);
2929
//! }
3030
//! ```
31+
//!
32+
//! `GenIterReturn` can be converted from a `Generator<()>`,
33+
//! `&mut GenIterReturn` can be used as iterator.
34+
//! The return value of the generator can be got after the iterator is done.
35+
//!
36+
//! ```
37+
//! #![feature(generators)]
38+
//!
39+
//! use gen_iter::gen_iter_return;
3140
//!
41+
//! let mut g = gen_iter_return!({
42+
//! yield 1;
43+
//! yield 2;
44+
//! return "done";
45+
//! });
46+
//!
47+
//! for y in &mut g {
48+
//! println!("yield {}", y);
49+
//! }
50+
//! println!("generator is_done={}", g.is_done()); // true
51+
//! println!("generator returns {}", g.return_or_self().ok().unwrap()); // "done"
52+
//! ```
3253
3354
#![no_std]
3455
#![feature(generators, generator_trait)]
3556
// #![feature(conservative_impl_trait)]
3657

3758
mod gen_iter;
3859
pub use gen_iter::*;
60+
61+
mod gen_iter_return;
62+
pub use gen_iter_return::*;

0 commit comments

Comments
 (0)