Add flat_map_ok iterator adaptor

src/flat_map_ok.rs

@@ -0,0 +1,206 @@
+use crate::size_hint;
+use std::{
+    fmt,
+    iter::{DoubleEndedIterator, FusedIterator},
+};
+
+pub fn flat_map_ok<I, F, T, U, E>(iter: I, f: F) -> FlatMapOk<I, F, T, U, E>
+where
+    I: Iterator<Item = Result<T, E>>,
+    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<I, F, T, U, E>
+where
+    I: Iterator<Item = Result<T, E>>,
+    F: FnMut(T) -> U,
+    U: IntoIterator,
+{
+    iter: I,
+    f: F,
+    inner_front: Option<U::IntoIter>,
+    inner_back: Option<U::IntoIter>,
+    _phantom: std::marker::PhantomData<fn() -> (T, E)>,
+}
+
+impl<I, F, T, U, E> Iterator for FlatMapOk<I, F, T, U, E>
+where
+    I: Iterator<Item = Result<T, E>>,
+    F: FnMut(T) -> U,
+    U: IntoIterator,
+{
+    type Item = Result<U::Item, E>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        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<B, 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<usize>) {
+        let inner_hint = |inner: &Option<U::IntoIter>| {
+            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<I, F, T, U, E> DoubleEndedIterator for FlatMapOk<I, F, T, U, E>
+where
+    I: DoubleEndedIterator<Item = Result<T, E>>,
+    F: FnMut(T) -> U,
+    U: IntoIterator,
+    U::IntoIter: DoubleEndedIterator,
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        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<B, Fold>(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<I, F, T, U, E> Clone for FlatMapOk<I, F, T, U, E>
+where
+    I: Iterator<Item = Result<T, E>> + Clone,
+    F: FnMut(T) -> U + Clone,
+    U: IntoIterator,
+    U::IntoIter: Clone,
+{
+    clone_fields!(iter, f, inner_front, inner_back, _phantom);
+}
+
+impl<I, F, T, U, E> fmt::Debug for FlatMapOk<I, F, T, U, E>
+where
+    I: Iterator<Item = Result<T, E>> + 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<I, F, T, U, E> FusedIterator for FlatMapOk<I, F, T, U, E>
+where
+    I: FusedIterator<Item = Result<T, E>>,
+    F: FnMut(T) -> U,
+    U: IntoIterator,
+{
+}

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<F, T, U, E>(self, f: F) -> FlatMapOk<Self, F, T, U, E>
+    where
+        Self: Iterator<Item = Result<T, E>> + 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.
     ///

tests/flat_map_ok.rs

@@ -0,0 +1,85 @@
+use itertools::{assert_equal, Itertools};
+use std::vec::IntoIter;
+
+fn mix_data() -> IntoIter<Result<i32, bool>> {
+    vec![Ok(2), Err(false), Ok(3), Err(true), Ok(0)].into_iter()
+}
+
+fn ok_data() -> IntoIter<Result<i32, bool>> {
+    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::<Result<Vec<_>, _>>(),
+        Err(false)
+    );
+}
+
+#[test]
+fn flat_map_ok_collect_ok_forward() {
+    assert_eq!(
+        ok_data()
+            .flat_map_ok(|i| 0..i)
+            .collect::<Result<Vec<_>, _>>(),
+        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::<Result<Vec<_>, _>>(),
+        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<Result<i32, bool>> = 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<Result<i32, bool>> = 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<Result<i32, bool>> = 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);
+}