Rollup merge of #80670 - the8472:fix-zip-trusted-random-access-composition, r=m-ou-se

TrustedRandomAaccess specialization composes incorrectly for nested iter::Zips

I found this while working on improvements for TRA.

After partially consuming a Zip adapter and then wrapping it into another Zip where the adapters use their `TrustedRandomAccess` specializations leads to the outer adapter returning elements which should have already been consumed.

If the optimizer gets tripped up by the addition this might affect performance for chained `zip()` iterators even when the inner one is not partially advanced but it would require more extensive fixes to `TrustedRandomAccess` to communicate those offsets earlier.

Included test fails on nightly, [playground link](https://play.rust-lang.org/?version=nightly&mode=debug&edition=2018&gist=24fa1edf8a104ff31f5a24830593b01f)

3 files changed, 23 insertions, 0 deletions

diff --git a/library/core/src/iter/adapters/zip.rs b/library/core/src/iter/adapters/zip.rs
index 5766fd3c887..98b8dca9614 100644
--- a/library/core/src/iter/adapters/zip.rs
+++ b/library/core/src/iter/adapters/zip.rs
@@ -286,6 +286,7 @@ where
 
     #[inline]
     unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item {
+        let idx = self.index + idx;
         // SAFETY: the caller must uphold the contract for
         // `Iterator::__iterator_get_unchecked`.
         unsafe { (self.a.__iterator_get_unchecked(idx), self.b.__iterator_get_unchecked(idx)) }
diff --git a/library/core/tests/iter.rs b/library/core/tests/iter.rs
index 691767edea6..bc5421bfb5f 100644
--- a/library/core/tests/iter.rs
+++ b/library/core/tests/iter.rs
@@ -2,6 +2,7 @@
 
 use core::cell::Cell;
 use core::convert::TryFrom;
+use core::iter::TrustedRandomAccess;
 use core::iter::*;
 
 /// An iterator wrapper that panics whenever `next` or `next_back` is called
@@ -602,6 +603,26 @@ fn test_zip_nth_back_side_effects_exhausted() {
 }
 
 #[test]
+fn test_zip_trusted_random_access_composition() {
+    let a = [0, 1, 2, 3, 4];
+    let b = a;
+    let c = a;
+
+    let a = a.iter().copied();
+    let b = b.iter().copied();
+    let mut c = c.iter().copied();
+    c.next();
+
+    let mut z1 = a.zip(b);
+    assert_eq!(z1.next().unwrap(), (0, 0));
+
+    let mut z2 = z1.zip(c);
+    fn assert_trusted_random_access<T: TrustedRandomAccess>(_a: &T) {}
+    assert_trusted_random_access(&z2);
+    assert_eq!(z2.next().unwrap(), ((1, 1), 1));
+}
+
+#[test]
 fn test_iterator_step_by() {
     // Identity
     let mut it = (0..).step_by(1).take(3);
diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs
index bc737cd1927..98f5982fbb2 100644
--- a/library/core/tests/lib.rs
+++ b/library/core/tests/lib.rs
@@ -75,6 +75,7 @@
 #![feature(const_option)]
 #![feature(integer_atomics)]
 #![feature(slice_group_by)]
+#![feature(trusted_random_access)]
 #![deny(unsafe_op_in_unsafe_fn)]
 
 extern crate test;