diff --git a/library/core/src/iter/adapters/take.rs b/library/core/src/iter/adapters/take.rs index 297dd0acaddc1..4c8f9fe16da0f 100644 --- a/library/core/src/iter/adapters/take.rs +++ b/library/core/src/iter/adapters/take.rs @@ -317,3 +317,60 @@ impl SpecTake for Take { } } } + +#[stable(feature = "exact_size_take_repeat", since = "CURRENT_RUSTC_VERSION")] +impl DoubleEndedIterator for Take> { + #[inline] + fn next_back(&mut self) -> Option { + self.next() + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option { + self.nth(n) + } + + #[inline] + fn try_rfold(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try, + { + self.try_fold(init, fold) + } + + #[inline] + fn rfold(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> { + self.advance_by(n) + } +} + +// Note: It may be tempting to impl DoubleEndedIterator for Take. +// 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 ExactSizeIterator for Take> { + fn len(&self) -> usize { + self.n + } +} + +#[stable(feature = "exact_size_take_repeat", since = "CURRENT_RUSTC_VERSION")] +impl A, A> ExactSizeIterator for Take> { + 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 39afa2cbfcaf2..65a8a93b4a916 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)); + } +}