Rollup merge of #94115 - scottmcm:iter-process-by-ref, r=yaahc

Let `try_collect` take advantage of `try_fold` overrides

No public API changes.

With this change, `try_collect` (#94047) is no longer going through the `impl Iterator for &mut impl Iterator`, and thus will be able to use `try_fold` overrides instead of being forced through `next` for every element.

Here's the test added, to see that it fails before this PR (once a new enough nightly is out): https://play.rust-lang.org/?version=nightly&mode=debug&edition=2021&gist=462f2896f2fed2c238ee63ca1a7e7c56

This might as well go to the same person as my last `try_process` PR  (#93572), so
r? ``@yaahc``

5 files changed, 72 insertions, 2 deletions

diff --git a/library/core/src/iter/adapters/by_ref_sized.rs b/library/core/src/iter/adapters/by_ref_sized.rs
new file mode 100644
index 00000000000..0b5e2a89ef3
--- /dev/null
+++ b/library/core/src/iter/adapters/by_ref_sized.rs
@@ -0,0 +1,42 @@
+use crate::ops::Try;
+
+/// Like `Iterator::by_ref`, but requiring `Sized` so it can forward generics.
+///
+/// Ideally this will no longer be required, eventually, but as can be seen in
+/// the benchmarks (as of Feb 2022 at least) `by_ref` can have performance cost.
+pub(crate) struct ByRefSized<'a, I>(pub &'a mut I);
+
+impl<I: Iterator> Iterator for ByRefSized<'_, I> {
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.0.next()
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.0.size_hint()
+    }
+
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        self.0.advance_by(n)
+    }
+
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        self.0.nth(n)
+    }
+
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        self.0.fold(init, f)
+    }
+
+    fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Output = B>,
+    {
+        self.0.try_fold(init, f)
+    }
+}
diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs
index 2ae92e89d63..d82fde47520 100644
--- a/library/core/src/iter/adapters/mod.rs
+++ b/library/core/src/iter/adapters/mod.rs
@@ -1,6 +1,7 @@
 use crate::iter::{InPlaceIterable, Iterator};
 use crate::ops::{ChangeOutputType, ControlFlow, FromResidual, NeverShortCircuit, Residual, Try};
 
+mod by_ref_sized;
 mod chain;
 mod cloned;
 mod copied;
@@ -31,6 +32,8 @@ pub use self::{
     scan::Scan, skip::Skip, skip_while::SkipWhile, take::Take, take_while::TakeWhile, zip::Zip,
 };
 
+pub(crate) use self::by_ref_sized::ByRefSized;
+
 #[stable(feature = "iter_cloned", since = "1.1.0")]
 pub use self::cloned::Cloned;
 
diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs
index 5a987733134..145c5ee109d 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::try_process;
+pub(crate) use self::adapters::{try_process, ByRefSized};
 
 mod adapters;
 mod range;
diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs
index d980f593081..b62e8dfe1d6 100644
--- a/library/core/src/iter/traits/iterator.rs
+++ b/library/core/src/iter/traits/iterator.rs
@@ -2,6 +2,7 @@ use crate::cmp::{self, Ordering};
 use crate::ops::{ChangeOutputType, ControlFlow, FromResidual, Residual, Try};
 
 use super::super::try_process;
+use super::super::ByRefSized;
 use super::super::TrustedRandomAccessNoCoerce;
 use super::super::{Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse};
 use super::super::{FlatMap, Flatten};
@@ -1861,7 +1862,7 @@ pub trait Iterator {
         <<Self as Iterator>::Item as Try>::Residual: Residual<B>,
         B: FromIterator<<Self::Item as Try>::Output>,
     {
-        try_process(self, |i| i.collect())
+        try_process(ByRefSized(self), |i| i.collect())
     }
 
     /// Collects all the items from an iterator into a collection.
diff --git a/library/core/tests/iter/traits/iterator.rs b/library/core/tests/iter/traits/iterator.rs
index 32bd68e3d25..731b1592d41 100644
--- a/library/core/tests/iter/traits/iterator.rs
+++ b/library/core/tests/iter/traits/iterator.rs
@@ -551,6 +551,30 @@ fn test_collect_into() {
     assert!(a == b);
 }
 
+#[test]
+fn iter_try_collect_uses_try_fold_not_next() {
+    // This makes sure it picks up optimizations, and doesn't use the `&mut I` impl.
+    struct PanicOnNext<I>(I);
+    impl<I: Iterator> Iterator for PanicOnNext<I> {
+        type Item = I::Item;
+        fn next(&mut self) -> Option<Self::Item> {
+            panic!("Iterator::next should not be called!")
+        }
+        fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R
+        where
+            Self: Sized,
+            F: FnMut(B, Self::Item) -> R,
+            R: std::ops::Try<Output = B>,
+        {
+            self.0.try_fold(init, f)
+        }
+    }
+
+    let it = (0..10).map(Some);
+    let _ = PanicOnNext(it).try_collect::<Vec<_>>();
+    // validation is just that it didn't panic.
+}
+
 // just tests by whether or not this compiles
 fn _empty_impl_all_auto_traits<T>() {
     use std::panic::{RefUnwindSafe, UnwindSafe};