Auto merge of #93572 - scottmcm:generic-iter-process, r=yaahc

Change `ResultShunt` to be generic over `Try`

Just a refactor (and rename) for now, so it's not `Result`-specific.

This could be used for a future `Iterator::try_collect`, or similar, but anything like that is left for a future PR.

6 files changed, 54 insertions, 48 deletions

diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs
index db8776ac741..2ae92e89d63 100644
--- a/library/core/src/iter/adapters/mod.rs
+++ b/library/core/src/iter/adapters/mod.rs
@@ -1,5 +1,5 @@
 use crate::iter::{InPlaceIterable, Iterator};
-use crate::ops::{ControlFlow, Try};
+use crate::ops::{ChangeOutputType, ControlFlow, FromResidual, NeverShortCircuit, Residual, Try};
 
 mod chain;
 mod cloned;
@@ -128,41 +128,45 @@ pub unsafe trait SourceIter {
 }
 
 /// An iterator adapter that produces output as long as the underlying
-/// iterator produces `Result::Ok` values.
+/// iterator produces values where `Try::branch` says to `ControlFlow::Continue`.
 ///
-/// If an error is encountered, the iterator stops and the error is
-/// stored.
-pub(crate) struct ResultShunt<'a, I, E> {
+/// If a `ControlFlow::Break` is encountered, the iterator stops and the
+/// residual is stored.
+pub(crate) struct GenericShunt<'a, I, R> {
     iter: I,
-    error: &'a mut Result<(), E>,
+    residual: &'a mut Option<R>,
 }
 
-/// 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>
+/// Process the given iterator as if it yielded a the item's `Try::Output`
+/// type instead. Any `Try::Residual`s encountered will stop the inner iterator
+/// and be propagated back to the overall result.
+pub(crate) fn try_process<I, T, R, F, U>(iter: I, mut f: F) -> ChangeOutputType<I::Item, U>
 where
-    I: Iterator<Item = Result<T, E>>,
-    for<'a> F: FnMut(ResultShunt<'a, I, E>) -> U,
+    I: Iterator<Item: Try<Output = T, Residual = R>>,
+    for<'a> F: FnMut(GenericShunt<'a, I, R>) -> U,
+    R: Residual<U>,
 {
-    let mut error = Ok(());
-    let shunt = ResultShunt { iter, error: &mut error };
+    let mut residual = None;
+    let shunt = GenericShunt { iter, residual: &mut residual };
     let value = f(shunt);
-    error.map(|()| value)
+    match residual {
+        Some(r) => FromResidual::from_residual(r),
+        None => Try::from_output(value),
+    }
 }
 
-impl<I, T, E> Iterator for ResultShunt<'_, I, E>
+impl<I, R> Iterator for GenericShunt<'_, I, R>
 where
-    I: Iterator<Item = Result<T, E>>,
+    I: Iterator<Item: Try<Residual = R>>,
 {
-    type Item = T;
+    type Item = <I::Item as Try>::Output;
 
     fn next(&mut self) -> Option<Self::Item> {
-        self.find(|_| true)
+        self.try_for_each(ControlFlow::Break).break_value()
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        if self.error.is_err() {
+        if self.residual.is_some() {
             (0, Some(0))
         } else {
             let (_, upper) = self.iter.size_hint();
@@ -170,17 +174,16 @@ where
         }
     }
 
-    fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
+    fn try_fold<B, F, T>(&mut self, init: B, mut f: F) -> T
     where
-        F: FnMut(B, Self::Item) -> R,
-        R: Try<Output = B>,
+        F: FnMut(B, Self::Item) -> T,
+        T: Try<Output = B>,
     {
-        let error = &mut *self.error;
         self.iter
-            .try_fold(init, |acc, x| match x {
-                Ok(x) => ControlFlow::from_try(f(acc, x)),
-                Err(e) => {
-                    *error = Err(e);
+            .try_fold(init, |acc, x| match Try::branch(x) {
+                ControlFlow::Continue(x) => ControlFlow::from_try(f(acc, x)),
+                ControlFlow::Break(r) => {
+                    *self.residual = Some(r);
                     ControlFlow::Break(try { acc })
                 }
             })
@@ -192,17 +195,12 @@ where
         Self: Sized,
         F: FnMut(B, Self::Item) -> B,
     {
-        #[inline]
-        fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
-            move |acc, x| Ok(f(acc, x))
-        }
-
-        self.try_fold(init, ok(fold)).unwrap()
+        self.try_fold(init, NeverShortCircuit::wrap_mut_2(fold)).0
     }
 }
 
 #[unstable(issue = "none", feature = "inplace_iteration")]
-unsafe impl<I, E> SourceIter for ResultShunt<'_, I, E>
+unsafe impl<I, R> SourceIter for GenericShunt<'_, I, R>
 where
     I: SourceIter,
 {
@@ -215,11 +213,11 @@ where
     }
 }
 
-// SAFETY: ResultShunt::next calls I::find, which has to advance `iter` in order to
-// return `Some(_)`. Since `iter` has type `I: InPlaceIterable` it's guaranteed that
-// at least one item will be moved out from the underlying source.
+// SAFETY: GenericShunt::next calls `I::try_for_each`, which has to advance `iter`
+// in order to return `Some(_)`. Since `iter` has type `I: InPlaceIterable` it's
+// guaranteed that at least one item will be moved out from the underlying source.
 #[unstable(issue = "none", feature = "inplace_iteration")]
-unsafe impl<I, T, E> InPlaceIterable for ResultShunt<'_, I, E> where
-    I: Iterator<Item = Result<T, E>> + InPlaceIterable
+unsafe impl<I, T, R> InPlaceIterable for GenericShunt<'_, I, R> where
+    I: Iterator<Item: Try<Output = T, Residual = R>> + InPlaceIterable
 {
 }
diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs
index da459ed7c68..65f56f64dbf 100644
--- a/library/core/src/iter/mod.rs
+++ b/library/core/src/iter/mod.rs
@@ -417,7 +417,7 @@ pub use self::adapters::{
 #[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
 pub use self::adapters::{Intersperse, IntersperseWith};
 
-pub(crate) use self::adapters::process_results;
+pub(crate) use self::adapters::try_process;
 
 mod adapters;
 mod range;
diff --git a/library/core/src/iter/traits/accum.rs b/library/core/src/iter/traits/accum.rs
index c2e837df5ff..84d83ee3969 100644
--- a/library/core/src/iter/traits/accum.rs
+++ b/library/core/src/iter/traits/accum.rs
@@ -167,7 +167,7 @@ where
     where
         I: Iterator<Item = Result<U, E>>,
     {
-        iter::process_results(iter, |i| i.sum())
+        iter::try_process(iter, |i| i.sum())
     }
 }
 
@@ -183,7 +183,7 @@ where
     where
         I: Iterator<Item = Result<U, E>>,
     {
-        iter::process_results(iter, |i| i.product())
+        iter::try_process(iter, |i| i.product())
     }
 }
 
@@ -210,7 +210,7 @@ where
     where
         I: Iterator<Item = Option<U>>,
     {
-        iter.map(|x| x.ok_or(())).sum::<Result<_, _>>().ok()
+        iter::try_process(iter, |i| i.sum())
     }
 }
 
@@ -226,6 +226,6 @@ where
     where
         I: Iterator<Item = Option<U>>,
     {
-        iter.map(|x| x.ok_or(())).product::<Result<_, _>>().ok()
+        iter::try_process(iter, |i| i.product())
     }
 }
diff --git a/library/core/src/ops/try_trait.rs b/library/core/src/ops/try_trait.rs
index 6a414ae8c4b..eac426ad311 100644
--- a/library/core/src/ops/try_trait.rs
+++ b/library/core/src/ops/try_trait.rs
@@ -359,6 +359,14 @@ pub(crate) type ChangeOutputType<T, V> = <<T as Try>::Residual as Residual<V>>::
 #[repr(transparent)]
 pub(crate) struct NeverShortCircuit<T>(pub T);
 
+impl<T> NeverShortCircuit<T> {
+    /// Wrap a binary `FnMut` to return its result wrapped in a `NeverShortCircuit`.
+    #[inline]
+    pub fn wrap_mut_2<A, B>(mut f: impl FnMut(A, B) -> T) -> impl FnMut(A, B) -> Self {
+        move |a, b| NeverShortCircuit(f(a, b))
+    }
+}
+
 pub(crate) enum NeverShortCircuitResidual {}
 
 impl<T> Try for NeverShortCircuit<T> {
diff --git a/library/core/src/option.rs b/library/core/src/option.rs
index 611f4ab38ab..ec04692d3e0 100644
--- a/library/core/src/option.rs
+++ b/library/core/src/option.rs
@@ -500,7 +500,7 @@
 
 #![stable(feature = "rust1", since = "1.0.0")]
 
-use crate::iter::{FromIterator, FusedIterator, TrustedLen};
+use crate::iter::{self, FromIterator, FusedIterator, TrustedLen};
 use crate::panicking::{panic, panic_str};
 use crate::pin::Pin;
 use crate::{
@@ -2233,7 +2233,7 @@ 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.
 
-        iter.into_iter().map(|x| x.ok_or(())).collect::<Result<_, _>>().ok()
+        iter::try_process(iter.into_iter(), |i| i.collect())
     }
 }
 
diff --git a/library/core/src/result.rs b/library/core/src/result.rs
index fbd6d419236..05b4fa035b1 100644
--- a/library/core/src/result.rs
+++ b/library/core/src/result.rs
@@ -2016,7 +2016,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.
 
-        iter::process_results(iter.into_iter(), |i| i.collect())
+        iter::try_process(iter.into_iter(), |i| i.collect())
     }
 }