diff --git a/src/lib.rs b/src/lib.rs index c6518ff60..8fbee99ca 100644 --- a/src/lib.rs +++ b/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]); diff --git a/src/sources.rs b/src/sources.rs index 562801bf0..caa49edbd 100644 --- a/src/sources.rs +++ b/src/sources.rs @@ -135,6 +135,83 @@ impl Iterator for Unfold } } +/// 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(seed: T, step: F) -> Generate + where F: FnMut(&T) -> Option +{ + Generate::new(Some(seed), step) +} + +impl fmt::Debug for Generate + 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 { + next: Option, + step: F, +} + +impl Generate + where F: FnMut(&T) -> Option +{ + /// Like [`generate`](../fn.generate.html), but allows for + /// empty sequences. + pub fn new(seed: Option, step: F) -> Generate { + Generate { + next: seed, + step: step, + } + } +} + +impl Iterator for Generate + where F: FnMut(&T) -> Option +{ + type Item = T; + + #[inline] + fn next(&mut self) -> Option { + self.next.take().map(|next| { + self.next = (self.step)(&next); + next + }) + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (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. diff --git a/tests/test_core.rs b/tests/test_core.rs index e21525501..a8a401fc2 100644 --- a/tests/test_core.rs +++ b/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); +}