Auto merge of rust-lang#85528 - the8472:iter-markers, r=dtolnay

Implement iterator specialization traits on more adapters

This adds

* `TrustedLen` to `Skip` and `StepBy`
* `TrustedRandomAccess` to `Skip`
* `InPlaceIterable` and `SourceIter` to  `Copied` and `Cloned`

The first two might improve performance in the compiler itself since `skip` is used in several places. Constellations that would exercise the last point are probably rare since it would require an owning iterator that has references as Items somewhere in its iterator pipeline.

Improvements for `Skip`:

```
# old
test iter::bench_skip_trusted_random_access                     ... bench:       8,335 ns/iter (+/- 90)

# new
test iter::bench_skip_trusted_random_access                     ... bench:       2,753 ns/iter (+/- 27)
```

Author: bors

Diff for: library/alloc/tests/vec.rs

@@ -1166,10 +1166,14 @@ fn test_from_iter_partially_drained_in_place_specialization() {
 #[test]
 fn test_from_iter_specialization_with_iterator_adapters() {
     fn assert_in_place_trait<T: InPlaceIterable>(_: &T) {}
-    let src: Vec<usize> = vec![0usize; 256];
+    let owned: Vec<usize> = vec![0usize; 256];
+    let refd: Vec<&usize> = owned.iter().collect();
+    let src: Vec<&&usize> = refd.iter().collect();
     let srcptr = src.as_ptr();
     let iter = src
         .into_iter()
+        .copied()
+        .cloned()
         .enumerate()
         .map(|i| i.0 + i.1)
         .zip(std::iter::repeat(1usize))
@@ -1180,7 +1184,7 @@ fn test_from_iter_specialization_with_iterator_adapters() {
     assert_in_place_trait(&iter);
     let sink = iter.collect::<Result<Vec<_>, _>>().unwrap();
     let sinkptr = sink.as_ptr();
-    assert_eq!(srcptr, sinkptr as *const usize);
+    assert_eq!(srcptr as *const usize, sinkptr as *const usize);
 }
 
 #[test]

Diff for: library/core/benches/iter.rs

@@ -391,6 +391,19 @@ fn bench_skip_then_zip(b: &mut Bencher) {
     });
 }
 
+#[bench]
+fn bench_skip_trusted_random_access(b: &mut Bencher) {
+    let v: Vec<u64> = black_box(vec![42; 10000]);
+    let mut sink = [0; 10000];
+
+    b.iter(|| {
+        for (val, idx) in v.iter().skip(8).zip(0..10000) {
+            sink[idx] += val;
+        }
+        sink
+    });
+}
+
 #[bench]
 fn bench_filter_count(b: &mut Bencher) {
     b.iter(|| (0i64..1000000).map(black_box).filter(|x| x % 3 == 0).count())

Diff for: library/core/src/iter/adapters/cloned.rs

@@ -1,8 +1,9 @@
 use crate::iter::adapters::{
-    zip::try_get_unchecked, TrustedRandomAccess, TrustedRandomAccessNoCoerce,
+    zip::try_get_unchecked, SourceIter, TrustedRandomAccess, TrustedRandomAccessNoCoerce,
 };
-use crate::iter::{FusedIterator, TrustedLen, UncheckedIterator};
+use crate::iter::{FusedIterator, InPlaceIterable, TrustedLen, UncheckedIterator};
 use crate::ops::Try;
+use core::num::NonZeroUsize;
 
 /// An iterator that clones the elements of an underlying iterator.
 ///
@@ -167,3 +168,23 @@ impl<I: Default> Default for Cloned<I> {
         Self::new(Default::default())
     }
 }
+
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<I> SourceIter for Cloned<I>
+where
+    I: SourceIter,
+{
+    type Source = I::Source;
+
+    #[inline]
+    unsafe fn as_inner(&mut self) -> &mut I::Source {
+        // SAFETY: unsafe function forwarding to unsafe function with the same requirements
+        unsafe { SourceIter::as_inner(&mut self.it) }
+    }
+}
+
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<I: InPlaceIterable> InPlaceIterable for Cloned<I> {
+    const EXPAND_BY: Option<NonZeroUsize> = I::EXPAND_BY;
+    const MERGE_BY: Option<NonZeroUsize> = I::MERGE_BY;
+}

Diff for: library/core/src/iter/adapters/copied.rs

@@ -1,7 +1,7 @@
 use crate::iter::adapters::{
-    zip::try_get_unchecked, TrustedRandomAccess, TrustedRandomAccessNoCoerce,
+    zip::try_get_unchecked, SourceIter, TrustedRandomAccess, TrustedRandomAccessNoCoerce,
 };
-use crate::iter::{FusedIterator, TrustedLen};
+use crate::iter::{FusedIterator, InPlaceIterable, TrustedLen};
 use crate::mem::MaybeUninit;
 use crate::mem::SizedTypeProperties;
 use crate::num::NonZeroUsize;
@@ -255,3 +255,23 @@ impl<I: Default> Default for Copied<I> {
         Self::new(Default::default())
     }
 }
+
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<I> SourceIter for Copied<I>
+where
+    I: SourceIter,
+{
+    type Source = I::Source;
+
+    #[inline]
+    unsafe fn as_inner(&mut self) -> &mut I::Source {
+        // SAFETY: unsafe function forwarding to unsafe function with the same requirements
+        unsafe { SourceIter::as_inner(&mut self.it) }
+    }
+}
+
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<I: InPlaceIterable> InPlaceIterable for Copied<I> {
+    const EXPAND_BY: Option<NonZeroUsize> = I::EXPAND_BY;
+    const MERGE_BY: Option<NonZeroUsize> = I::MERGE_BY;
+}

Diff for: library/core/src/iter/adapters/skip.rs

@@ -1,6 +1,10 @@
 use crate::intrinsics::unlikely;
+use crate::iter::adapters::zip::try_get_unchecked;
 use crate::iter::TrustedFused;
-use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable};
+use crate::iter::{
+    adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen, TrustedRandomAccess,
+    TrustedRandomAccessNoCoerce,
+};
 use crate::num::NonZeroUsize;
 use crate::ops::{ControlFlow, Try};
 
@@ -152,6 +156,32 @@ where
 
         NonZeroUsize::new(n).map_or(Ok(()), Err)
     }
+
+    #[doc(hidden)]
+    unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item
+    where
+        Self: TrustedRandomAccessNoCoerce,
+    {
+        // SAFETY: the caller must uphold the contract for
+        // `Iterator::__iterator_get_unchecked`.
+        //
+        // Dropping the skipped prefix when index 0 is passed is safe
+        // since
+        // * the caller passing index 0 means that the inner iterator has more items than `self.n`
+        // * TRA contract requires that get_unchecked will only be called once
+        //   (unless elements are copyable)
+        // * it does not conflict with in-place iteration since index 0 must be accessed
+        //   before something is written into the storage used by the prefix
+        unsafe {
+            if Self::MAY_HAVE_SIDE_EFFECT && idx == 0 {
+                for skipped_idx in 0..self.n {
+                    drop(try_get_unchecked(&mut self.iter, skipped_idx));
+                }
+            }
+
+            try_get_unchecked(&mut self.iter, idx + self.n)
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -237,3 +267,23 @@ unsafe impl<I: InPlaceIterable> InPlaceIterable for Skip<I> {
     const EXPAND_BY: Option<NonZeroUsize> = I::EXPAND_BY;
     const MERGE_BY: Option<NonZeroUsize> = I::MERGE_BY;
 }
+
+#[doc(hidden)]
+#[unstable(feature = "trusted_random_access", issue = "none")]
+unsafe impl<I> TrustedRandomAccess for Skip<I> where I: TrustedRandomAccess {}
+
+#[doc(hidden)]
+#[unstable(feature = "trusted_random_access", issue = "none")]
+unsafe impl<I> TrustedRandomAccessNoCoerce for Skip<I>
+where
+    I: TrustedRandomAccessNoCoerce,
+{
+    const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT;
+}
+
+// SAFETY: This adapter is shortening. TrustedLen requires the upper bound to be calculated correctly.
+// These requirements can only be satisfied when the upper bound of the inner iterator's upper
+// bound is never `None`. I: TrustedRandomAccess happens to provide this guarantee while
+// I: TrustedLen would not.
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<I> TrustedLen for Skip<I> where I: Iterator + TrustedRandomAccess {}

Diff for: library/core/src/iter/adapters/step_by.rs

@@ -1,7 +1,7 @@
 use crate::convert::TryFrom;
 use crate::{
     intrinsics,
-    iter::{from_fn, TrustedLen},
+    iter::{from_fn, TrustedLen, TrustedRandomAccess},
     ops::{Range, Try},
 };
 
@@ -124,6 +124,14 @@ where
 #[stable(feature = "iterator_step_by", since = "1.28.0")]
 impl<I> ExactSizeIterator for StepBy<I> where I: ExactSizeIterator {}
 
+// SAFETY: This adapter is shortening. TrustedLen requires the upper bound to be calculated correctly.
+// These requirements can only be satisfied when the upper bound of the inner iterator's upper
+// bound is never `None`. I: TrustedRandomAccess happens to provide this guarantee while
+// I: TrustedLen would not.
+// This also covers the Range specializations since the ranges also implement TRA
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<I> TrustedLen for StepBy<I> where I: Iterator + TrustedRandomAccess {}
+
 trait SpecRangeSetup<T> {
     fn setup(inner: T, step: usize) -> T;
 }
@@ -480,12 +488,6 @@ macro_rules! spec_int_ranges {
                 acc
             }
         }
-
-        /// Safety: This macro is only applied to ranges over types <= usize
-        /// which means the inner length is guaranteed to fit into a usize and so
-        /// the outer length calculation won't encounter clamped values
-        #[unstable(feature = "trusted_len", issue = "37572")]
-        unsafe impl TrustedLen for StepBy<Range<$t>> {}
     )*)
 }
 

Diff for: library/core/tests/iter/adapters/step_by.rs

@@ -220,7 +220,8 @@ fn test_iterator_step_by_size_hint() {
     assert_eq!(it.len(), 3);
 
     // Cannot be TrustedLen as a step greater than one makes an iterator
-    // with (usize::MAX, None) no longer meet the safety requirements
+    // with (usize::MAX, None) no longer meet the safety requirements.
+    // Exception: The inner iterator is known to have a len() <= usize::MAX
     trait TrustedLenCheck {
         fn test(self) -> bool;
     }
@@ -235,7 +236,9 @@ fn test_iterator_step_by_size_hint() {
         }
     }
     assert!(TrustedLenCheck::test(a.iter()));
-    assert!(!TrustedLenCheck::test(a.iter().step_by(1)));
+    assert!(TrustedLenCheck::test(a.iter().step_by(1)));
+    assert!(TrustedLenCheck::test(a.iter().chain(a.iter())));
+    assert!(!TrustedLenCheck::test(a.iter().chain(a.iter()).step_by(1)));
 }
 
 #[test]

Diff for: library/core/tests/iter/range.rs

@@ -474,6 +474,16 @@ fn test_range_inclusive_size_hint() {
     assert_eq!((imin..=imax + 1).size_hint(), (usize::MAX, None));
 }
 
+#[test]
+fn test_range_trusted_random_access() {
+    let mut range = 0..10;
+    unsafe {
+        assert_eq!(range.next(), Some(0));
+        assert_eq!(range.__iterator_get_unchecked(0), 1);
+        assert_eq!(range.__iterator_get_unchecked(1), 2);
+    }
+}
+
 #[test]
 fn test_double_ended_range() {
     assert_eq!((11..14).rev().collect::<Vec<_>>(), [13, 12, 11]);