Merge pull request #6 from viruscamp/gen_iter_return

add struct `GenIterReturn` and macro `gen_iter_return!` to iterate over a generator and get the return value

Author: tinaun

src/gen_iter.rs

@@ -0,0 +1,106 @@
+use core::ops::{Generator, GeneratorState};
+use core::iter::Iterator;
+use core::marker::Unpin;
+use core::pin::Pin;
+
+/// an iterator that holds an internal generator representing
+/// the iteration state
+#[derive(Copy, Clone, Debug)]
+pub struct GenIter<T>(pub T)
+where
+    T: Generator<Return = ()> + Unpin;
+
+impl<T> Iterator for GenIter<T>
+where
+    T: Generator<Return = ()> + Unpin,
+{
+    type Item = T::Yield;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        match Pin::new(&mut self.0).resume(()) {
+            GeneratorState::Yielded(n) => Some(n),
+            GeneratorState::Complete(()) => None,
+        }
+    }
+}
+
+impl<G> From<G> for GenIter<G>
+where
+    G: Generator<Return = ()> + Unpin,
+{
+    #[inline]
+    fn from(gen: G) -> Self {
+        GenIter(gen)
+    }
+}
+
+
+/// macro to simplify iterator - via - generator construction
+///
+/// ```
+/// #![feature(generators)]
+///
+/// use gen_iter::gen_iter;
+///
+/// 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);
+///
+/// ```
+#[macro_export]
+macro_rules! gen_iter {
+    ($block: block) => {
+        $crate::GenIter(|| $block)
+    };
+    (move $block: block) => {
+        $crate::GenIter(move || $block)
+    }
+}
+
+
+#[cfg(test)]
+mod tests {
+    use super::GenIter;
+
+    #[test]
+    fn it_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 into_gen_iter() {
+        let mut g: GenIter<_> = (|| {
+            yield 1;
+            yield 2;
+        }).into();
+
+        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);
+    }
+}

src/gen_iter_return.rs

@@ -0,0 +1,163 @@
+use core::ops::{Generator, GeneratorState};
+use core::iter::{Iterator, FusedIterator};
+use core::marker::Unpin;
+use core::pin::Pin;
+
+/// `GenIterReturn<G>` holds a generator `G` or the return value of `G`,
+/// `&mut GenIterReturn<G>` acts as an iterator.
+/// 
+/// Differences with `GenIter<G>`:
+/// 1. able to get return value of a generator
+/// 2. safe to call `next()` after generator is done without panic
+/// 3. maybe less efficient than `GenIter<G>`
+#[derive(Copy, Clone, Debug)]
+pub struct GenIterReturn<G: Generator + Unpin>(Result<G::Return, G>);
+
+impl<G: Generator + 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),
+        }
+    }
+}
+
+/// Force use `&mut g` as iterator to prevent the code below,
+/// in which return value cannot be got.
+/// ```compile_fail
+/// // !!INVALID CODE!!
+/// # #![feature(generators)]
+/// # use gen_iter::gen_iter_return;
+/// let mut g = gen_iter_return!({ yield 1; return "done"; });
+/// for v in g {} // invalid, because `GenIterReturn<G>` is not `Iterator`
+/// let ret = g.return_or_self(); // g is dropped after for loop
+/// ```
+impl<G: Generator + 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(()) {
+                GeneratorState::Yielded(y) => Some(y),
+                GeneratorState::Complete(r) => {
+                    self.0 = Ok(r);
+                    None
+                },
+            }
+        }
+    }
+}
+
+/// `GenIterReturn<G>` satisfies the trait `FusedIterator`
+impl<G: Generator + Unpin> FusedIterator for &mut GenIterReturn<G> {}
+
+impl<G: Generator + Unpin> From<G> for GenIterReturn<G> {
+    #[inline]
+    fn from(g: G) -> Self {
+        GenIterReturn::new(g)
+    }
+}
+
+/// macro to simplify iterator - via - generator with return value construction
+/// ```
+/// #![feature(generators)]
+///
+/// use gen_iter::gen_iter_return;
+///
+/// let mut g = gen_iter_return!({
+///     yield 1;
+///     yield 2;
+///     return "done";
+/// });
+///
+/// assert_eq!((&mut g).collect::<Vec<_>>(), [1, 2]); // use `&mut g` as an iterator
+/// assert_eq!(g.is_done(), true); // check whether generator is done
+/// assert_eq!((&mut g).next(), None); // safe to call `next()` after done
+/// assert_eq!(g.return_or_self().ok(), Some("done")); // get return value of generator
+/// ```
+#[macro_export]
+macro_rules! gen_iter_return {
+    ($block: block) => {
+        $crate::GenIterReturn::new(|| $block)
+    };
+    (move $block: block) => {
+        $crate::GenIterReturn::new(move || $block)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::GenIterReturn;
+
+    /// test `new` and all instance method,
+    /// and show that it won't panic when call `next()` even exhausted.
+    #[test]
+    fn it_works() {
+        let mut g = GenIterReturn::new(|| {
+            yield 1;
+            return "done";
+        });
+
+        assert_eq!((&mut g).next(), Some(1));
+        assert_eq!(g.is_done(), false);
+
+        g = match g.return_or_self() {
+            Ok(_) => panic!("generator is done but should not"),
+            Err(g) => g
+        };
+
+        assert_eq!((&mut g).next(), None);
+        assert_eq!(g.is_done(), true);
+
+        assert_eq!((&mut g).next(), None); // it won't panic when call `next()` even exhausted.
+
+        assert_eq!(g.return_or_self().ok(), Some("done"));
+    }
+
+    #[test]
+    fn from_generator() {
+        let mut g = GenIterReturn::from(|| {
+            yield 1;
+            return "done";
+        });
+
+        assert_eq!((&mut g).next(), Some(1));
+        assert_eq!((&mut g).next(), None);
+
+        assert_eq!(g.is_done(), true);
+        assert_eq!(g.return_or_self().ok(), Some("done"));
+    }
+
+    /// normal usage using macro `gen_iter_return`
+    #[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_eq!(g.is_done(), true);
+        assert_eq!(g.return_or_self().ok(), Some("done"));
+    }
+}