implement TrustedLen for StepBy

3 files changed, 24 insertions, 9 deletions

diff --git a/library/core/src/iter/adapters/step_by.rs b/library/core/src/iter/adapters/step_by.rs
index 9e83584e3f3..54ed4c952fb 100644
--- a/library/core/src/iter/adapters/step_by.rs
+++ b/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 --git a/library/core/tests/iter/adapters/step_by.rs b/library/core/tests/iter/adapters/step_by.rs
index 4c5b1dd9a6b..67360e38131 100644
--- a/library/core/tests/iter/adapters/step_by.rs
+++ b/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 --git a/library/core/tests/iter/range.rs b/library/core/tests/iter/range.rs
index 5b87d6c1fa0..a6b9f1cb7c8 100644
--- a/library/core/tests/iter/range.rs
+++ b/library/core/tests/iter/range.rs
@@ -475,6 +475,16 @@ fn test_range_inclusive_size_hint() {
 }
 
 #[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]);
     for _ in (10..0).rev() {