From a2a7caacf7e8e1ec78ac23e01aa281312ec780c2 Mon Sep 17 00:00:00 2001
From: The8472 <git@infinite-source.de>
Date: Thu, 20 May 2021 00:29:20 +0200
Subject: [PATCH] implement TrustedLen for StepBy

---
 library/core/src/iter/adapters/step_by.rs   | 16 +++++++++-------
 library/core/tests/iter/adapters/step_by.rs |  7 +++++--
 library/core/tests/iter/range.rs            | 10 ++++++++++
 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 @@ fn rfold<Acc, F>(self, init: Acc, f: F) -> Acc
 #[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 @@ fn spec_fold<Acc, F>(self, init: Acc, mut f: F) -> Acc
                 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 size_hint(&self) -> (usize, Option<usize>) {
     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(self) -> bool {
         }
     }
     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
@@ -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]);