Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 57 additions & 0 deletions library/core/src/iter/adapters/take.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
90 changes: 90 additions & 0 deletions library/core/tests/iter/adapters/take.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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));
}
}