Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement DoubleEnded and ExactSize for Take<Repeat> and Take<RepeatWith> #106943

Merged
merged 1 commit into from
Aug 17, 2024
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>> {
mina86 marked this conversation as resolved.
Show resolved Hide resolved
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));
}
}
Loading