Repository: rust-itertools/itertools (language: rust)

Issue:
flatten_ok
## Proposal

I propose adding the following (without the `impl Iterator`, just to demonstrate):

```rust
    fn flatten_ok<A, B, E>(self) -> impl Iterator<Item=Result<B, E>>
        where Self: Iterator<Item = Result<A, E>>, A: IntoIterator<Item = B>
    {...}
```

`flatten_ok` operates on an iterator of `Result<A, E>` where `A: IntoIterator<Item = B>`. If the iterator produces `Ok`, the value is flattened. If the iterator produces `Err`, the value is produced exactly as-is.

## Why add this?

Several others have been trying to solve this exact problem:
* https://paulkernfeld.com/2018/11/03/flatten-nested-iterator-result.html
* https://stackoverflow.com/questions/59852161/how-to-handle-result-in-flat-map

I also encountered this problem and was unable to find a combinator for this purpose, which was surprising to me.

## Potential way to implement

It could be implemented this way:
```rust
fn flatten_ok<A, B, E>(self) -> impl Iterator<Item=Result<B, E>>
where Self: Iterator<Item = Result<A, E>>, A: IntoIterator<Item = B>
{
    self.flat_map(|res| Either::from(res)
        .map_right(|ok| ok.into_iter().map(Ok))
        .map_left(Err)
        .map_left(std::iter::once)
    )
}
```

Of course, a struct could be created within `itertools` to implement the `Iterator` so there is no external API dependency on `either`, and so we can still specify a concrete return type, but the above would approximately the implementation used and does work if this signature (with `impl Iterator`) is acceptable.

## Questions

* Is there another and more clear way to implement this that others are missing? If so, please let me know.
* Should this be proposed to `libcore` first?


Discussion / hints:
If this proposal (with or without `impl Iterator`) is acceptable, I can send in a PR. Just let me know if its good.
Yes, this proposal seems acceptable! Note:
- Unfortunately, using impl trait isn't an option on trait methods (which is what Itertools provides). You'll need to define a concrete `FlattenOk` return type.
- We already depend on `Either`, so no issue there!

You're welcome to propose this to libcore first, but you don't have to. I'm happy to accept a PR for this here. (The absolute worst case scenario is that we might, someday, need to remove it from Itertools if it gets upstreamed into libcore. That's fine — we *do* make breaking changes from time to time anyways — and the underlying cause for this breakage may soon be addressed by a language RFC.)
@jswrenn Excellent! I will get to work on this later today after work.
I would be wary of using `Either` with iterators since it adds an additional discriminant check for each `next` call. 

In my opinion a faster implementation would be taking [`FlattenCompat`](https://github.com/rust-lang/rust/blob/master/library/core/src/iter/adapters/flatten.rs#L233-L431) from the stdlib and special casing it for `Result`. In particular when `iter` yields an `Err` I would not store any iterator in `frontiter` or `backiter` because that iterator would already be exhausted.
I've faced the same issue. For me, a solution without `Either` was based on `scan`: https://users.rust-lang.org/t/handling-errors-inside-flat-map-iterator/55238/11
@mpizenberg I am going to start working on a PR for this today. I will take a look at the scan version too. It sounds like, for efficiency reasons, it may be best to just do it from scratch though, but I can compare.

Fix patch:
diff --git a/src/flatten_ok.rs b/src/flatten_ok.rs
new file mode 100644
index 000000000..d46bbde4e
--- /dev/null
+++ b/src/flatten_ok.rs
@@ -0,0 +1,166 @@
+use crate::size_hint;
+use std::{
+    fmt,
+    iter::{DoubleEndedIterator, FusedIterator},
+};
+
+pub fn flatten_ok<I, T, E>(iter: I) -> FlattenOk<I, T, E>
+where
+    I: Iterator<Item = Result<T, E>>,
+    T: IntoIterator,
+{
+    FlattenOk {
+        iter,
+        inner_front: None,
+        inner_back: None,
+    }
+}
+
+/// An iterator adaptor that flattens `Result::Ok` values and
+/// allows `Result::Err` values through unchanged.
+///
+/// See [`.flatten_ok()`](crate::Itertools::flatten_ok) for more information.
+#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
+pub struct FlattenOk<I, T, E>
+where
+    I: Iterator<Item = Result<T, E>>,
+    T: IntoIterator,
+{
+    iter: I,
+    inner_front: Option<T::IntoIter>,
+    inner_back: Option<T::IntoIter>,
+}
+
+impl<I, T, E> Iterator for FlattenOk<I, T, E>
+where
+    I: Iterator<Item = Result<T, E>>,
+    T: IntoIterator,
+{
+    type Item = Result<T::Item, E>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            // Handle the front inner iterator.
+            if let Some(inner) = &mut self.inner_front {
+                if let Some(item) = inner.next() {
+                    return Some(Ok(item));
+                } else {
+                    // This is necessary for the iterator to implement `FusedIterator`
+                    // with only the orginal iterator being fused.
+                    self.inner_front = None;
+                }
+            }
+
+            match self.iter.next() {
+                Some(Ok(ok)) => self.inner_front = Some(ok.into_iter()),
+                Some(Err(e)) => return Some(Err(e)),
+                None => {
+                    // Handle the back inner iterator.
+                    if let Some(inner) = &mut self.inner_back {
+                        if let Some(item) = inner.next() {
+                            return Some(Ok(item));
+                        } else {
+                            // This is necessary for the iterator to implement `FusedIterator`
+                            // with only the orginal iterator being fused.
+                            self.inner_back = None;
+                        }
+                    } else {
+                        return None;
+                    }
+                }
+            }
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let inner_hint = |inner: &Option<T::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);
+        // The outer iterator `Ok` case could be (0, None) as we don't know its size_hint yet.
+        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, T, E> DoubleEndedIterator for FlattenOk<I, T, E>
+where
+    I: DoubleEndedIterator<Item = Result<T, E>>,
+    T: IntoIterator,
+    T::IntoIter: DoubleEndedIterator,
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        loop {
+            // Handle the back inner iterator.
+            if let Some(inner) = &mut self.inner_back {
+                if let Some(item) = inner.next_back() {
+                    return Some(Ok(item));
+                } else {
+                    // This is necessary for the iterator to implement `FusedIterator`
+                    // with only the orginal iterator being fused.
+                    self.inner_back = None;
+                }
+            }
+
+            match self.iter.next_back() {
+                Some(Ok(ok)) => self.inner_back = Some(ok.into_iter()),
+                Some(Err(e)) => return Some(Err(e)),
+                None => {
+                    // Handle the front inner iterator.
+                    if let Some(inner) = &mut self.inner_front {
+                        if let Some(item) = inner.next_back() {
+                            return Some(Ok(item));
+                        } else {
+                            // This is necessary for the iterator to implement `FusedIterator`
+                            // with only the orginal iterator being fused.
+                            self.inner_front = None;
+                        }
+                    } else {
+                        return None;
+                    }
+                }
+            }
+        }
+    }
+}
+
+impl<I, T, E> Clone for FlattenOk<I, T, E>
+where
+    I: Iterator<Item = Result<T, E>> + Clone,
+    T: IntoIterator,
+    T::IntoIter: Clone,
+{
+    #[inline]
+    clone_fields!(iter, inner_front, inner_back);
+}
+
+impl<I, T, E> fmt::Debug for FlattenOk<I, T, E>
+where
+    I: Iterator<Item = Result<T, E>> + fmt::Debug,
+    T: IntoIterator,
+    T::IntoIter: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("FlattenOk")
+            .field("iter", &self.iter)
+            .field("inner_front", &self.inner_front)
+            .field("inner_back", &self.inner_back)
+            .finish()
+    }
+}
+
+/// Only the iterator being flattened needs to implement [`FusedIterator`].
+impl<I, T, E> FusedIterator for FlattenOk<I, T, E>
+where
+    I: FusedIterator<Item = Result<T, E>>,
+    T: IntoIterator,
+{
+}
diff --git a/src/lib.rs b/src/lib.rs
index 4db8f5d40..cdf45137a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -119,6 +119,7 @@ pub mod structs {
     pub use crate::cons_tuples_impl::ConsTuples;
     pub use crate::exactly_one_err::ExactlyOneError;
     pub use crate::format::{Format, FormatWith};
+    pub use crate::flatten_ok::FlattenOk;
     #[cfg(feature = "use_std")]
     pub use crate::grouping_map::{GroupingMap, GroupingMapBy};
     #[cfg(feature = "use_alloc")]
@@ -194,6 +195,7 @@ mod combinations;
 mod combinations_with_replacement;
 mod exactly_one_err;
 mod diff;
+mod flatten_ok;
 mod format;
 #[cfg(feature = "use_std")]
 mod grouping_map;
@@ -833,6 +835,30 @@ pub trait Itertools : Iterator {
         adaptors::filter_map_ok(self, f)
     }
 
+    /// Return an iterator adaptor that flattens every `Result::Ok` value into
+    /// a series of `Result::Ok` values. `Result::Err` values are unchanged.
+    /// 
+    /// This is useful when you have some common error type for your crate and
+    /// need to propogate it upwards, but the `Result::Ok` case needs to be flattened.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// let input = vec![Ok(0..2), Err(false), Ok(2..4)];
+    /// let it = input.iter().cloned().flatten_ok();
+    /// itertools::assert_equal(it.clone(), vec![Ok(0), Ok(1), Err(false), Ok(2), Ok(3)]);
+    /// 
+    /// // This can also be used to propogate errors when collecting.
+    /// let output_result: Result<Vec<i32>, bool> = it.collect();
+    /// assert_eq!(output_result, Err(false));
+    /// ```
+    fn flatten_ok<T, E>(self) -> FlattenOk<Self, T, E>
+        where Self: Iterator<Item = Result<T, E>> + Sized,
+              T: IntoIterator
+    {
+        flatten_ok::flatten_ok(self)
+    }
+
     /// Return an iterator adaptor that merges the two base iterators in
     /// ascending order.  If both base iterators are sorted (ascending), the
     /// result is sorted.


Test patch:
diff --git a/tests/flatten_ok.rs b/tests/flatten_ok.rs
new file mode 100644
index 000000000..bf835b5d7
--- /dev/null
+++ b/tests/flatten_ok.rs
@@ -0,0 +1,76 @@
+use itertools::{assert_equal, Itertools};
+use std::{ops::Range, vec::IntoIter};
+
+fn mix_data() -> IntoIter<Result<Range<i32>, bool>> {
+    vec![Ok(0..2), Err(false), Ok(2..4), Err(true), Ok(4..6)].into_iter()
+}
+
+fn ok_data() -> IntoIter<Result<Range<i32>, bool>> {
+    vec![Ok(0..2), Ok(2..4), Ok(4..6)].into_iter()
+}
+
+#[test]
+fn flatten_ok_mixed_expected_forward() {
+    assert_equal(
+        mix_data().flatten_ok(),
+        vec![
+            Ok(0),
+            Ok(1),
+            Err(false),
+            Ok(2),
+            Ok(3),
+            Err(true),
+            Ok(4),
+            Ok(5),
+        ],
+    );
+}
+
+#[test]
+fn flatten_ok_mixed_expected_reverse() {
+    assert_equal(
+        mix_data().flatten_ok().rev(),
+        vec![
+            Ok(5),
+            Ok(4),
+            Err(true),
+            Ok(3),
+            Ok(2),
+            Err(false),
+            Ok(1),
+            Ok(0),
+        ],
+    );
+}
+
+#[test]
+fn flatten_ok_collect_mixed_forward() {
+    assert_eq!(
+        mix_data().flatten_ok().collect::<Result<Vec<_>, _>>(),
+        Err(false)
+    );
+}
+
+#[test]
+fn flatten_ok_collect_mixed_reverse() {
+    assert_eq!(
+        mix_data().flatten_ok().rev().collect::<Result<Vec<_>, _>>(),
+        Err(true)
+    );
+}
+
+#[test]
+fn flatten_ok_collect_ok_forward() {
+    assert_eq!(
+        ok_data().flatten_ok().collect::<Result<Vec<_>, _>>(),
+        Ok((0..6).collect())
+    );
+}
+
+#[test]
+fn flatten_ok_collect_ok_reverse() {
+    assert_eq!(
+        ok_data().flatten_ok().rev().collect::<Result<Vec<_>, _>>(),
+        Ok((0..6).rev().collect())
+    );
+}