diff --git a/src/flat_map_ok.rs b/src/flat_map_ok.rs new file mode 100644 index 000000000..c89ef1623 --- /dev/null +++ b/src/flat_map_ok.rs @@ -0,0 +1,206 @@ +use crate::size_hint; +use std::{ + fmt, + iter::{DoubleEndedIterator, FusedIterator}, +}; + +pub fn flat_map_ok(iter: I, f: F) -> FlatMapOk +where + I: Iterator>, + F: FnMut(T) -> U, + U: IntoIterator, +{ + FlatMapOk { + iter, + f, + inner_front: None, + inner_back: None, + _phantom: std::marker::PhantomData, + } +} + +/// An iterator adaptor that applies a function to `Result::Ok` values and +/// flattens the resulting iterator. `Result::Err` values are passed through +/// unchanged. +/// +/// This is equivalent to `.map_ok(f).flatten_ok()`. +/// +/// See [`.flat_map_ok()`](crate::Itertools::flat_map_ok) for more information. +#[must_use = "iterator adaptors are lazy and do nothing unless consumed"] +pub struct FlatMapOk +where + I: Iterator>, + F: FnMut(T) -> U, + U: IntoIterator, +{ + iter: I, + f: F, + inner_front: Option, + inner_back: Option, + _phantom: std::marker::PhantomData (T, E)>, +} + +impl Iterator for FlatMapOk +where + I: Iterator>, + F: FnMut(T) -> U, + U: IntoIterator, +{ + type Item = Result; + + fn next(&mut self) -> Option { + loop { + if let Some(inner) = &mut self.inner_front { + if let Some(item) = inner.next() { + return Some(Ok(item)); + } + self.inner_front = None; + } + + match self.iter.next() { + Some(Ok(ok)) => self.inner_front = Some((self.f)(ok).into_iter()), + Some(Err(e)) => return Some(Err(e)), + None => { + if let Some(inner) = &mut self.inner_back { + if let Some(item) = inner.next() { + return Some(Ok(item)); + } + self.inner_back = None; + } else { + return None; + } + } + } + } + } + + fn fold(self, init: B, mut fold_f: Fold) -> B + where + Self: Sized, + Fold: FnMut(B, Self::Item) -> B, + { + let mut f = self.f; + + // Front + let mut acc = match self.inner_front { + Some(x) => x.fold(init, |a, o| fold_f(a, Ok(o))), + None => init, + }; + + acc = self.iter.fold(acc, |acc, x| match x { + Ok(ok) => f(ok).into_iter().fold(acc, |a, o| fold_f(a, Ok(o))), + Err(e) => fold_f(acc, Err(e)), + }); + + // Back + match self.inner_back { + Some(x) => x.fold(acc, |a, o| fold_f(a, Ok(o))), + None => acc, + } + } + + fn size_hint(&self) -> (usize, Option) { + let inner_hint = |inner: &Option| { + inner + .as_ref() + .map(Iterator::size_hint) + .unwrap_or((0, Some(0))) + }; + let inner_front = inner_hint(&self.inner_front); + let inner_back = inner_hint(&self.inner_back); + let outer = match self.iter.size_hint() { + (0, Some(0)) => (0, Some(0)), + _ => (0, None), + }; + + size_hint::add(size_hint::add(inner_front, inner_back), outer) + } +} + +impl DoubleEndedIterator for FlatMapOk +where + I: DoubleEndedIterator>, + F: FnMut(T) -> U, + U: IntoIterator, + U::IntoIter: DoubleEndedIterator, +{ + fn next_back(&mut self) -> Option { + loop { + if let Some(inner) = &mut self.inner_back { + if let Some(item) = inner.next_back() { + return Some(Ok(item)); + } + self.inner_back = None; + } + + match self.iter.next_back() { + Some(Ok(ok)) => self.inner_back = Some((self.f)(ok).into_iter()), + Some(Err(e)) => return Some(Err(e)), + None => { + if let Some(inner) = &mut self.inner_front { + if let Some(item) = inner.next_back() { + return Some(Ok(item)); + } + self.inner_front = None; + } else { + return None; + } + } + } + } + } + + fn rfold(self, init: B, mut fold_f: Fold) -> B + where + Self: Sized, + Fold: FnMut(B, Self::Item) -> B, + { + let mut f = self.f; + + // Back + let mut acc = match self.inner_back { + Some(x) => x.rfold(init, |a, o| fold_f(a, Ok(o))), + None => init, + }; + + acc = self.iter.rfold(acc, |acc, x| match x { + Ok(ok) => f(ok).into_iter().rfold(acc, |a, o| fold_f(a, Ok(o))), + Err(e) => fold_f(acc, Err(e)), + }); + + // Front + match self.inner_front { + Some(x) => x.rfold(acc, |a, o| fold_f(a, Ok(o))), + None => acc, + } + } +} + +impl Clone for FlatMapOk +where + I: Iterator> + Clone, + F: FnMut(T) -> U + Clone, + U: IntoIterator, + U::IntoIter: Clone, +{ + clone_fields!(iter, f, inner_front, inner_back, _phantom); +} + +impl fmt::Debug for FlatMapOk +where + I: Iterator> + fmt::Debug, + F: FnMut(T) -> U, + U: IntoIterator, + U::IntoIter: fmt::Debug, +{ + debug_fmt_fields!(FlatMapOk, iter, inner_front, inner_back); +} + +/// Only the iterator being flat-mapped needs to implement [`FusedIterator`]. +impl FusedIterator for FlatMapOk +where + I: FusedIterator>, + F: FnMut(T) -> U, + U: IntoIterator, +{ +} diff --git a/src/lib.rs b/src/lib.rs index a3fb9b1db..2e2172d9e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -106,6 +106,7 @@ pub mod structs { #[cfg(feature = "use_std")] pub use crate::duplicates_impl::{Duplicates, DuplicatesBy}; pub use crate::exactly_one_err::ExactlyOneError; + pub use crate::flat_map_ok::FlatMapOk; pub use crate::flatten_ok::FlattenOk; pub use crate::format::{Format, FormatWith}; #[allow(deprecated)] @@ -194,6 +195,7 @@ mod duplicates_impl; mod exactly_one_err; #[cfg(feature = "use_alloc")] mod extrema_set; +mod flat_map_ok; mod flatten_ok; mod format; #[cfg(feature = "use_alloc")] @@ -1148,6 +1150,28 @@ pub trait Itertools: Iterator { flatten_ok::flatten_ok(self) } + /// Return an iterator adaptor that applies a function to every `Result::Ok` + /// value and flattens the resulting iterator. `Result::Err` values are + /// unchanged. + /// + /// This is equivalent to `.map_ok(f).flatten_ok()`. + /// + /// ``` + /// use itertools::Itertools; + /// + /// let input = vec![Ok(0i32), Err(false), Ok(3i32)]; + /// let it = input.into_iter().flat_map_ok(|i| 0..i); + /// itertools::assert_equal(it, vec![Err(false), Ok(0), Ok(1), Ok(2)]); + /// ``` + fn flat_map_ok(self, f: F) -> FlatMapOk + where + Self: Iterator> + Sized, + F: FnMut(T) -> U, + U: IntoIterator, + { + flat_map_ok::flat_map_ok(self, f) + } + /// “Lift” a function of the values of the current iterator so as to process /// an iterator of `Result` values instead. /// diff --git a/tests/flat_map_ok.rs b/tests/flat_map_ok.rs new file mode 100644 index 000000000..3afb78800 --- /dev/null +++ b/tests/flat_map_ok.rs @@ -0,0 +1,85 @@ +use itertools::{assert_equal, Itertools}; +use std::vec::IntoIter; + +fn mix_data() -> IntoIter> { + vec![Ok(2), Err(false), Ok(3), Err(true), Ok(0)].into_iter() +} + +fn ok_data() -> IntoIter> { + vec![Ok(2), Ok(3), Ok(0)].into_iter() +} + +#[test] +fn flat_map_ok_mixed_forward() { + assert_equal( + mix_data().flat_map_ok(|i| 0..i), + vec![Ok(0), Ok(1), Err(false), Ok(0), Ok(1), Ok(2), Err(true)], + ); +} + +#[test] +fn flat_map_ok_mixed_reverse() { + assert_equal( + mix_data().flat_map_ok(|i| 0..i).rev(), + vec![Err(true), Ok(2), Ok(1), Ok(0), Err(false), Ok(1), Ok(0)], + ); +} + +#[test] +fn flat_map_ok_collect_mixed() { + assert_eq!( + mix_data() + .flat_map_ok(|i| 0..i) + .collect::, _>>(), + Err(false) + ); +} + +#[test] +fn flat_map_ok_collect_ok_forward() { + assert_eq!( + ok_data() + .flat_map_ok(|i| 0..i) + .collect::, _>>(), + Ok(vec![0, 1, 0, 1, 2]) + ); +} + +#[test] +fn flat_map_ok_collect_ok_reverse() { + assert_eq!( + ok_data() + .flat_map_ok(|i| 0..i) + .rev() + .collect::, _>>(), + Ok(vec![2, 1, 0, 1, 0]) + ); +} + +#[test] +fn flat_map_ok_empty_results() { + // When the mapping function returns an empty iterator for some Ok values + let data: Vec> = vec![Ok(0), Ok(2), Ok(0)]; + assert_equal(data.into_iter().flat_map_ok(|i| 0..i), vec![Ok(0), Ok(1)]); +} + +#[test] +fn flat_map_ok_all_errors() { + let data: Vec> = vec![Err(false), Err(true)]; + assert_equal( + data.into_iter().flat_map_ok(|i: i32| 0..i), + vec![Err(false), Err(true)], + ); +} + +#[test] +fn flat_map_ok_equivalence_with_map_ok_flatten_ok() { + // flat_map_ok(f) should be equivalent to map_ok(f).flatten_ok() + let data1: Vec> = vec![Ok(2), Err(false), Ok(3)]; + let data2 = data1.clone(); + + let result1: Vec<_> = data1.into_iter().flat_map_ok(|i| 0..i).collect(); + let result2: Vec<_> = data2.into_iter().map_ok(|i| 0..i).flatten_ok().collect(); + + assert_eq!(result1, result2); +}