Rollup merge of #62459 - timvermeulen:result_sum_internal_iteration, r=scottmcm

Use internal iteration in the Sum and Product impls of Result and Option

This PR adds internal iteration to the `ResultShunt` iterator type underlying the `Sum` and `Product` impls of `Result`. I had to change `ResultShunt` to hold a mutable reference to an error instead, similar to `itertools::ProcessResults`, in order to be able to pass the `ResultShunt` itself by value (which is necessary for internal iteration).

`ResultShunt::process` can unfortunately no longer be an associated function because that would make it generic over the lifetime of the error reference, which wouldn't work, so I turned it into the free function `process_results`.

I removed the `OptionShunt` type and forwarded the `Sum` and `Product` impls of `Option` to their respective impls of `Result` instead, to avoid having to repeat the internal iteration logic.

6 files changed, 67 insertions, 119 deletions

diff --git a/src/libcore/iter/adapters/mod.rs b/src/libcore/iter/adapters/mod.rs
index af46e6df294..b2702902956 100644
--- a/src/libcore/iter/adapters/mod.rs
+++ b/src/libcore/iter/adapters/mod.rs
@@ -2181,136 +2181,64 @@ impl<I: FusedIterator, F> FusedIterator for Inspect<I, F>
     where F: FnMut(&I::Item) {}
 
 /// An iterator adapter that produces output as long as the underlying
-/// iterator produces `Option::Some` values.
-pub(crate) struct OptionShunt<I> {
-    iter: I,
-    exited_early: bool,
-}
-
-impl<I, T> OptionShunt<I>
-where
-    I: Iterator<Item = Option<T>>,
-{
-    /// Process the given iterator as if it yielded a `T` instead of a
-    /// `Option<T>`. Any `None` value will stop the inner iterator and
-    /// the overall result will be a `None`.
-    pub fn process<F, U>(iter: I, mut f: F) -> Option<U>
-    where
-        F: FnMut(&mut Self) -> U,
-    {
-        let mut shunt = OptionShunt::new(iter);
-        let value = f(shunt.by_ref());
-        shunt.reconstruct(value)
-    }
-
-    fn new(iter: I) -> Self {
-        OptionShunt {
-            iter,
-            exited_early: false,
-        }
-    }
-
-    /// Consume the adapter and rebuild a `Option` value.
-    fn reconstruct<U>(self, val: U) -> Option<U> {
-        if self.exited_early {
-            None
-        } else {
-            Some(val)
-        }
-    }
-}
-
-impl<I, T> Iterator for OptionShunt<I>
-where
-    I: Iterator<Item = Option<T>>,
-{
-    type Item = T;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        match self.iter.next() {
-            Some(Some(v)) => Some(v),
-            Some(None) => {
-                self.exited_early = true;
-                None
-            }
-            None => None,
-        }
-    }
-
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        if self.exited_early {
-            (0, Some(0))
-        } else {
-            let (_, upper) = self.iter.size_hint();
-            (0, upper)
-        }
-    }
-}
-
-/// An iterator adapter that produces output as long as the underlying
 /// iterator produces `Result::Ok` values.
 ///
 /// If an error is encountered, the iterator stops and the error is
-/// stored. The error may be recovered later via `reconstruct`.
-pub(crate) struct ResultShunt<I, E> {
+/// stored.
+pub(crate) struct ResultShunt<'a, I, E> {
     iter: I,
-    error: Option<E>,
+    error: &'a mut Result<(), E>,
 }
 
-impl<I, T, E> ResultShunt<I, E>
-    where I: Iterator<Item = Result<T, E>>
+/// Process the given iterator as if it yielded a `T` instead of a
+/// `Result<T, _>`. Any errors will stop the inner iterator and
+/// the overall result will be an error.
+pub(crate) fn process_results<I, T, E, F, U>(iter: I, mut f: F) -> Result<U, E>
+where
+    I: Iterator<Item = Result<T, E>>,
+    for<'a> F: FnMut(ResultShunt<'a, I, E>) -> U,
 {
-    /// Process the given iterator as if it yielded a `T` instead of a
-    /// `Result<T, _>`. Any errors will stop the inner iterator and
-    /// the overall result will be an error.
-    pub fn process<F, U>(iter: I, mut f: F) -> Result<U, E>
-        where F: FnMut(&mut Self) -> U
-    {
-        let mut shunt = ResultShunt::new(iter);
-        let value = f(shunt.by_ref());
-        shunt.reconstruct(value)
-    }
-
-    fn new(iter: I) -> Self {
-        ResultShunt {
-            iter,
-            error: None,
-        }
-    }
-
-    /// Consume the adapter and rebuild a `Result` value. This should
-    /// *always* be called, otherwise any potential error would be
-    /// lost.
-    fn reconstruct<U>(self, val: U) -> Result<U, E> {
-        match self.error {
-            None => Ok(val),
-            Some(e) => Err(e),
-        }
-    }
+    let mut error = Ok(());
+    let shunt = ResultShunt {
+        iter,
+        error: &mut error,
+    };
+    let value = f(shunt);
+    error.map(|()| value)
 }
 
-impl<I, T, E> Iterator for ResultShunt<I, E>
+impl<I, T, E> Iterator for ResultShunt<'_, I, E>
     where I: Iterator<Item = Result<T, E>>
 {
     type Item = T;
 
     fn next(&mut self) -> Option<Self::Item> {
-        match self.iter.next() {
-            Some(Ok(v)) => Some(v),
-            Some(Err(e)) => {
-                self.error = Some(e);
-                None
-            }
-            None => None,
-        }
+        self.find(|_| true)
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        if self.error.is_some() {
+        if self.error.is_err() {
             (0, Some(0))
         } else {
             let (_, upper) = self.iter.size_hint();
             (0, upper)
         }
     }
+
+    fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Ok = B>,
+    {
+        let error = &mut *self.error;
+        self.iter
+            .try_fold(init, |acc, x| match x {
+                Ok(x) => LoopState::from_try(f(acc, x)),
+                Err(e) => {
+                    *error = Err(e);
+                    LoopState::Break(Try::from_ok(acc))
+                }
+            })
+            .into_try()
+    }
 }
diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs
index 4a7d7f96b9b..aba8e84d58b 100644
--- a/src/libcore/iter/mod.rs
+++ b/src/libcore/iter/mod.rs
@@ -360,7 +360,7 @@ pub use self::adapters::Flatten;
 #[stable(feature = "iter_copied", since = "1.36.0")]
 pub use self::adapters::Copied;
 
-pub(crate) use self::adapters::{TrustedRandomAccess, OptionShunt, ResultShunt};
+pub(crate) use self::adapters::{TrustedRandomAccess, process_results};
 
 mod range;
 mod sources;
diff --git a/src/libcore/iter/traits/accum.rs b/src/libcore/iter/traits/accum.rs
index 01b64fb08ac..812463e77f9 100644
--- a/src/libcore/iter/traits/accum.rs
+++ b/src/libcore/iter/traits/accum.rs
@@ -1,6 +1,6 @@
 use crate::ops::{Mul, Add};
 use crate::num::Wrapping;
-use crate::iter::adapters::{OptionShunt, ResultShunt};
+use crate::iter;
 
 /// Trait to represent types that can be created by summing up an iterator.
 ///
@@ -139,7 +139,7 @@ impl<T, U, E> Sum<Result<U, E>> for Result<T, E>
     fn sum<I>(iter: I) -> Result<T, E>
         where I: Iterator<Item = Result<U, E>>,
     {
-        ResultShunt::process(iter, |i| i.sum())
+        iter::process_results(iter, |i| i.sum())
     }
 }
 
@@ -153,7 +153,7 @@ impl<T, U, E> Product<Result<U, E>> for Result<T, E>
     fn product<I>(iter: I) -> Result<T, E>
         where I: Iterator<Item = Result<U, E>>,
     {
-        ResultShunt::process(iter, |i| i.product())
+        iter::process_results(iter, |i| i.product())
     }
 }
 
@@ -180,7 +180,7 @@ where
     where
         I: Iterator<Item = Option<U>>,
     {
-        OptionShunt::process(iter, |i| i.sum())
+        iter.map(|x| x.ok_or(())).sum::<Result<_, _>>().ok()
     }
 }
 
@@ -196,6 +196,6 @@ where
     where
         I: Iterator<Item = Option<U>>,
     {
-        OptionShunt::process(iter, |i| i.product())
+        iter.map(|x| x.ok_or(())).product::<Result<_, _>>().ok()
     }
 }
diff --git a/src/libcore/option.rs b/src/libcore/option.rs
index 7713e5761d4..259ed36c578 100644
--- a/src/libcore/option.rs
+++ b/src/libcore/option.rs
@@ -135,7 +135,7 @@
 
 #![stable(feature = "rust1", since = "1.0.0")]
 
-use crate::iter::{FromIterator, FusedIterator, TrustedLen, OptionShunt};
+use crate::iter::{FromIterator, FusedIterator, TrustedLen};
 use crate::{convert, fmt, hint, mem, ops::{self, Deref, DerefMut}};
 use crate::pin::Pin;
 
@@ -1499,7 +1499,10 @@ impl<A, V: FromIterator<A>> FromIterator<Option<A>> for Option<V> {
         // FIXME(#11084): This could be replaced with Iterator::scan when this
         // performance bug is closed.
 
-        OptionShunt::process(iter.into_iter(), |i| i.collect())
+        iter.into_iter()
+            .map(|x| x.ok_or(()))
+            .collect::<Result<_, _>>()
+            .ok()
     }
 }
 
diff --git a/src/libcore/result.rs b/src/libcore/result.rs
index 559877ddd5a..8c60a9c1b50 100644
--- a/src/libcore/result.rs
+++ b/src/libcore/result.rs
@@ -231,7 +231,7 @@
 #![stable(feature = "rust1", since = "1.0.0")]
 
 use crate::fmt;
-use crate::iter::{FromIterator, FusedIterator, TrustedLen, ResultShunt};
+use crate::iter::{self, FromIterator, FusedIterator, TrustedLen};
 use crate::ops::{self, Deref, DerefMut};
 
 /// `Result` is a type that represents either success ([`Ok`]) or failure ([`Err`]).
@@ -1343,7 +1343,7 @@ impl<A, E, V: FromIterator<A>> FromIterator<Result<A, E>> for Result<V, E> {
         // FIXME(#11084): This could be replaced with Iterator::scan when this
         // performance bug is closed.
 
-        ResultShunt::process(iter.into_iter(), |i| i.collect())
+        iter::process_results(iter.into_iter(), |i| i.collect())
     }
 }
 
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index e27e1605607..3615fab7915 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -1203,6 +1203,23 @@ fn test_iterator_sum_result() {
     assert_eq!(v.iter().cloned().sum::<Result<i32, _>>(), Ok(10));
     let v: &[Result<i32, ()>] = &[Ok(1), Err(()), Ok(3), Ok(4)];
     assert_eq!(v.iter().cloned().sum::<Result<i32, _>>(), Err(()));
+
+    #[derive(PartialEq, Debug)]
+    struct S(Result<i32, ()>);
+
+    impl Sum<Result<i32, ()>> for S {
+        fn sum<I: Iterator<Item = Result<i32, ()>>>(mut iter: I) -> Self {
+            // takes the sum by repeatedly calling `next` on `iter`,
+            // thus testing that repeated calls to `ResultShunt::try_fold`
+            // produce the expected results
+            Self(iter.by_ref().sum())
+        }
+    }
+
+    let v: &[Result<i32, ()>] = &[Ok(1), Ok(2), Ok(3), Ok(4)];
+    assert_eq!(v.iter().cloned().sum::<S>(), S(Ok(10)));
+    let v: &[Result<i32, ()>] = &[Ok(1), Err(()), Ok(3), Ok(4)];
+    assert_eq!(v.iter().cloned().sum::<S>(), S(Err(())));
 }
 
 #[test]