about summary refs log tree commit diff
diff options
context:
space:
mode:
-rw-r--r--library/core/src/iter/adapters/take.rs57
-rw-r--r--library/core/tests/iter/adapters/take.rs90
2 files changed, 147 insertions, 0 deletions
diff --git a/library/core/src/iter/adapters/take.rs b/library/core/src/iter/adapters/take.rs
index 297dd0acadd..4c8f9fe16da 100644
--- a/library/core/src/iter/adapters/take.rs
+++ b/library/core/src/iter/adapters/take.rs
@@ -317,3 +317,60 @@ impl<I: Iterator + TrustedRandomAccess> SpecTake for Take<I> {
         }
     }
 }
+
+#[stable(feature = "exact_size_take_repeat", since = "CURRENT_RUSTC_VERSION")]
+impl<T: Clone> DoubleEndedIterator for Take<crate::iter::Repeat<T>> {
+    #[inline]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.next()
+    }
+
+    #[inline]
+    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
+        self.nth(n)
+    }
+
+    #[inline]
+    fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
+    where
+        Self: Sized,
+        Fold: FnMut(Acc, Self::Item) -> R,
+        R: Try<Output = Acc>,
+    {
+        self.try_fold(init, fold)
+    }
+
+    #[inline]
+    fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
+    where
+        Self: Sized,
+        Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        self.fold(init, fold)
+    }
+
+    #[inline]
+    #[rustc_inherit_overflow_checks]
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
+        self.advance_by(n)
+    }
+}
+
+// Note: It may be tempting to impl DoubleEndedIterator for Take<RepeatWith>.
+// One must fight that temptation since such implementation wouldn’t be correct
+// because we have no way to return value of nth invocation of repeater followed
+// by n-1st without remembering all results.
+
+#[stable(feature = "exact_size_take_repeat", since = "CURRENT_RUSTC_VERSION")]
+impl<T: Clone> ExactSizeIterator for Take<crate::iter::Repeat<T>> {
+    fn len(&self) -> usize {
+        self.n
+    }
+}
+
+#[stable(feature = "exact_size_take_repeat", since = "CURRENT_RUSTC_VERSION")]
+impl<F: FnMut() -> A, A> ExactSizeIterator for Take<crate::iter::RepeatWith<F>> {
+    fn len(&self) -> usize {
+        self.n
+    }
+}
diff --git a/library/core/tests/iter/adapters/take.rs b/library/core/tests/iter/adapters/take.rs
index 39afa2cbfca..65a8a93b4a9 100644
--- a/library/core/tests/iter/adapters/take.rs
+++ b/library/core/tests/iter/adapters/take.rs
@@ -170,3 +170,93 @@ fn test_byref_take_consumed_items() {
     assert_eq!(count, 70);
     assert_eq!(inner, 90..90);
 }
+
+#[test]
+fn test_exact_size_take_repeat() {
+    let mut iter = core::iter::repeat(42).take(40);
+    assert_eq!((40, Some(40)), iter.size_hint());
+    assert_eq!(40, iter.len());
+
+    assert_eq!(Some(42), iter.next());
+    assert_eq!((39, Some(39)), iter.size_hint());
+    assert_eq!(39, iter.len());
+
+    assert_eq!(Some(42), iter.next_back());
+    assert_eq!((38, Some(38)), iter.size_hint());
+    assert_eq!(38, iter.len());
+
+    assert_eq!(Some(42), iter.nth(3));
+    assert_eq!((34, Some(34)), iter.size_hint());
+    assert_eq!(34, iter.len());
+
+    assert_eq!(Some(42), iter.nth_back(3));
+    assert_eq!((30, Some(30)), iter.size_hint());
+    assert_eq!(30, iter.len());
+
+    assert_eq!(Ok(()), iter.advance_by(10));
+    assert_eq!((20, Some(20)), iter.size_hint());
+    assert_eq!(20, iter.len());
+
+    assert_eq!(Ok(()), iter.advance_back_by(10));
+    assert_eq!((10, Some(10)), iter.size_hint());
+    assert_eq!(10, iter.len());
+}
+
+#[test]
+fn test_exact_size_take_repeat_with() {
+    let mut counter = 0;
+    let mut iter = core::iter::repeat_with(move || {
+        counter += 1;
+        counter
+    })
+    .take(40);
+    assert_eq!((40, Some(40)), iter.size_hint());
+    assert_eq!(40, iter.len());
+
+    assert_eq!(Some(1), iter.next());
+    assert_eq!((39, Some(39)), iter.size_hint());
+    assert_eq!(39, iter.len());
+
+    assert_eq!(Some(5), iter.nth(3));
+    assert_eq!((35, Some(35)), iter.size_hint());
+    assert_eq!(35, iter.len());
+
+    assert_eq!(Ok(()), iter.advance_by(10));
+    assert_eq!((25, Some(25)), iter.size_hint());
+    assert_eq!(25, iter.len());
+
+    assert_eq!(Some(16), iter.next());
+    assert_eq!((24, Some(24)), iter.size_hint());
+    assert_eq!(24, iter.len());
+}
+
+// This is https://github.com/rust-lang/rust/issues/104729 with all uses of
+// repeat(0) were replaced by repeat(0).take(20).
+#[test]
+fn test_reverse_on_zip() {
+    let vec_1 = [1; 10];
+
+    let zipped_iter = vec_1.iter().copied().zip(core::iter::repeat(0).take(20));
+
+    // Forward
+    for (one, zero) in zipped_iter {
+        assert_eq!((1, 0), (one, zero));
+    }
+
+    let rev_vec_iter = vec_1.iter().rev();
+    let rev_repeat_iter = std::iter::repeat(0).take(20).rev();
+
+    // Manual reversed zip
+    let rev_zipped_iter = rev_vec_iter.zip(rev_repeat_iter);
+
+    for (&one, zero) in rev_zipped_iter {
+        assert_eq!((1, 0), (one, zero));
+    }
+
+    let zipped_iter = vec_1.iter().zip(core::iter::repeat(0).take(20));
+
+    // Cannot call rev here for automatic reversed zip constuction
+    for (&one, zero) in zipped_iter.rev() {
+        assert_eq!((1, 0), (one, zero));
+    }
+}