Add `and_then_ok` and `and_then_err`

2 files changed, 104 insertions, 0 deletions

diff --git a/src/and_then_x.rs b/src/and_then_x.rs
new file mode 100644
index 0000000..3173fa2
--- /dev/null
+++ b/src/and_then_x.rs
@@ -0,0 +1,103 @@
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+//
+
+/// Extension trait for `Iterator<Item = Result<O, E>>` to selectively transform Oks and Errors.
+pub trait AndThenX<O, E>: Sized {
+    fn and_then_ok<F>(self, F) -> AndThenOk<Self, F>
+    where
+        F: FnMut(O) -> Result<O, E>;
+    fn and_then_err<F>(self, F) -> AndThenErr<Self, F>
+    where
+        F: FnMut(E) -> Result<O, E>;
+}
+
+impl<I, O, E> AndThenX<O, E> for I
+where
+    I: Iterator<Item = Result<O, E>> + Sized,
+{
+    fn and_then_ok<F>(self, f: F) -> AndThenOk<Self, F>
+    where
+        F: FnMut(O) -> Result<O, E>,
+    {
+        AndThenOk { iter: self, f }
+    }
+    fn and_then_err<F>(self, f: F) -> AndThenErr<Self, F>
+    where
+        F: FnMut(E) -> Result<O, E>,
+    {
+        AndThenErr { iter: self, f }
+    }
+}
+
+#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
+pub struct AndThenOk<I, F> {
+    iter: I,
+    f: F,
+}
+
+impl<I, O, E, F> Iterator for AndThenOk<I, F>
+where
+    I: Iterator<Item = Result<O, E>>,
+    F: FnMut(O) -> Result<O, E>,
+{
+    type Item = Result<O, E>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.iter.next() {
+            Some(Ok(o)) => Some((self.f)(o)),
+            other => other,
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
+}
+
+#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
+pub struct AndThenErr<I, F> {
+    iter: I,
+    f: F,
+}
+
+impl<I, O, E, F> Iterator for AndThenErr<I, F>
+where
+    I: Iterator<Item = Result<O, E>>,
+    F: FnMut(E) -> Result<O, E>,
+{
+    type Item = Result<O, E>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.iter.next() {
+            Some(Err(e)) => Some((self.f)(e)),
+            other => other,
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
+}
+
+#[test]
+fn test_and_then_ok() {
+    use std::str::FromStr;
+
+    let mapped: Vec<_> = ["1", "2", "a", "4", "5"]
+        .into_iter()
+        .map(|txt| usize::from_str(txt))
+        .and_then_ok(|i| Ok(2 * i))
+        .and_then_err(|i| Ok(15))
+        .collect();
+
+    assert_eq!(mapped[0], Ok(2));
+    assert_eq!(mapped[1], Ok(4));
+    assert_eq!(mapped[2], Ok(15));
+    assert_eq!(mapped[3], Ok(8));
+    assert_eq!(mapped[4], Ok(10));
+}
diff --git a/src/lib.rs b/src/lib.rs
index 6ad6fc9..471f8fe 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -194,6 +194,7 @@
 //! MPL 2.0
 //!
 
+pub mod and_then_x;
 pub mod errors;
 pub mod filter_map_x;
 pub mod filter_x;