Add itertools::generate

src/lib.rs

@@ -91,7 +91,7 @@ pub mod structs {
     #[cfg(feature = "use_std")]
     pub use rciter_impl::RcIter;
     pub use repeatn::RepeatN;
-    pub use sources::{RepeatCall, Unfold, Iterate};
+    pub use sources::{RepeatCall, Generate, Unfold, Iterate};
     #[cfg(feature = "use_std")]
     pub use tee::Tee;
     pub use tuple_impl::{TupleBuffer, TupleWindows, Tuples};
@@ -114,7 +114,7 @@ pub use minmax::MinMaxResult;
 pub use peeking_take_while::PeekingNext;
 pub use process_results_impl::process_results;
 pub use repeatn::repeat_n;
-pub use sources::{repeat_call, unfold, iterate};
+pub use sources::{repeat_call, generate, unfold, iterate};
 pub use with_position::Position;
 pub use ziptuple::multizip;
 mod adaptors;
@@ -1888,13 +1888,13 @@ pub trait Itertools : Iterator {
 
     /// Return a `HashMap` of keys mapped to `Vec`s of values. Keys and values
     /// are taken from `(Key, Value)` tuple pairs yielded by the input iterator.
-    /// 
+    ///
     /// ```
     /// use itertools::Itertools;
-    /// 
+    ///
     /// let data = vec![(0, 10), (2, 12), (3, 13), (0, 20), (3, 33), (2, 42)];
     /// let lookup = data.into_iter().into_group_map();
-    /// 
+    ///
     /// assert_eq!(lookup[&0], vec![10, 20]);
     /// assert_eq!(lookup.get(&1), None);
     /// assert_eq!(lookup[&2], vec![12, 42]);

src/sources.rs

@@ -135,6 +135,83 @@ impl<A, St, F> Iterator for Unfold<St, F>
     }
 }
 
+/// Returns an iterator that starts with the `seed` and repeatedly
+/// calls the `step` function to produce the next value, until it
+/// returns `None`. The resulting sequence always has at least one
+/// element -- the `seed`. To produce a possibly empty sequence,
+/// use `Generate::new`.
+///
+/// ```
+/// // an iterator that yields Collatz sequence.
+/// // https://en.wikipedia.org/wiki/Collatz_conjecture
+///
+/// use itertools::generate;
+///
+/// let collatz = generate(12, |&n| {
+///     if n == 1 {
+///         None
+///     } else if n % 2 == 0 {
+///         Some(n / 2)
+///     } else {
+///         Some(3 * n + 1)
+///     }
+/// });
+///
+/// itertools::assert_equal(collatz,
+///                         vec![12, 6, 3, 10, 5, 16, 8, 4, 2, 1]);
+/// ```
+pub fn generate<T, F>(seed: T, step: F) -> Generate<T, F>
+    where F: FnMut(&T) -> Option<T>
+{
+    Generate::new(Some(seed), step)
+}
+
+impl<T, F> fmt::Debug for Generate<T, F>
+    where T: fmt::Debug,
+{
+    debug_fmt_fields!(Genrate, next);
+}
+
+/// See [`generate`](../fn.generate.html) for more information.
+#[derive(Clone)]
+#[must_use = "iterators are lazy and do nothing unless consumed"]
+pub struct Generate<T, F> {
+    next: Option<T>,
+    step: F,
+}
+
+impl<T, F> Generate<T, F>
+    where F: FnMut(&T) -> Option<T>
+{
+    /// Like [`generate`](../fn.generate.html), but allows for
+    /// empty sequences.
+    pub fn new(seed: Option<T>, step: F) -> Generate<T, F> {
+        Generate {
+            next: seed,
+            step: step,
+        }
+    }
+}
+
+impl<T, F> Iterator for Generate<T, F>
+    where F: FnMut(&T) -> Option<T>
+{
+    type Item = T;
+
+    #[inline]
+    fn next(&mut self) -> Option<T> {
+        self.next.take().map(|next| {
+            self.next = (self.step)(&next);
+            next
+        })
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.next.is_some() as usize, None)
+    }
+}
+
 /// An iterator that infinitely applies function to value and yields results.
 ///
 /// This `struct` is created by the [`iterate()`] function. See its documentation for more.

tests/test_core.rs

@@ -12,6 +12,7 @@ use it::Itertools;
 use it::interleave;
 use it::multizip;
 use it::free::put_back;
+use it::{generate, Generate};
 
 #[test]
 fn product2() {
@@ -238,3 +239,13 @@ fn tree_fold1() {
         assert_eq!((0..i).tree_fold1(|x, y| x + y), (0..i).fold1(|x, y| x + y));
     }
 }
+
+#[test]
+fn test_generate() {
+    let count_down = generate(5u32, |&n| n.checked_sub(1));
+    it::assert_equal(count_down, [5, 4, 3, 2, 1, 0].iter().cloned());
+    let non_empty = generate((), |&()| None);
+    assert_eq!(non_empty.count(), 1);
+    let empty = Generate::new(None, |&()| None);
+    assert_eq!(empty.count(), 0);
+}