Implement iterator specialization traits on more adapters

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]

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())

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;
+}

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;
+}

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 {}

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>> {}
     )*)
 }
 

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]

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]);