diff --git a/library/core/src/iter/adapters/take.rs b/library/core/src/iter/adapters/take.rs index ce18bffe7146f..bdc2a92fb4a20 100644 --- a/library/core/src/iter/adapters/take.rs +++ b/library/core/src/iter/adapters/take.rs @@ -1,5 +1,7 @@ use crate::cmp; -use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen}; +use crate::iter::{ + adapters::SourceIter, FusedIterator, InPlaceIterable, Repeat, RepeatWith, TrustedLen, +}; use crate::num::NonZeroUsize; use crate::ops::{ControlFlow, Try}; @@ -241,9 +243,34 @@ where } } +#[stable(feature = "double_ended_repeat_take_iterator", since = "CURRENT_RUSTC_VERSION")] +impl DoubleEndedIterator for Take> { + fn next_back(&mut self) -> Option { + if self.n == 0 { + None + } else { + self.n -= 1; + self.iter.next() + } + } + + fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> { + let advanced = self.n.min(n); + self.n -= advanced; + let remainder = n - advanced; + NonZeroUsize::new(remainder).map_or(Ok(()), Err) + } +} + #[stable(feature = "rust1", since = "1.0.0")] impl ExactSizeIterator for Take where I: ExactSizeIterator {} +#[stable(feature = "take_repeat_exactsize", since = "CURRENT_RUSTC_VERSION")] +impl ExactSizeIterator for Take> {} + +#[stable(feature = "take_repeat_with_exactsize", since = "CURRENT_RUSTC_VERSION")] +impl ExactSizeIterator for Take> where F: FnMut() -> T {} + #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Take where I: FusedIterator {} diff --git a/library/core/tests/iter/adapters/take.rs b/library/core/tests/iter/adapters/take.rs index 3cad47c06de03..d9c4532fef44f 100644 --- a/library/core/tests/iter/adapters/take.rs +++ b/library/core/tests/iter/adapters/take.rs @@ -167,3 +167,22 @@ fn test_byref_take_consumed_items() { assert_eq!(count, 70); assert_eq!(inner, 90..90); } + +#[test] +fn test_take_repeat_double_ended() { + let mut iter = repeat(0).take(100); + assert_eq!(iter.len(), 100); + assert_eq!(iter.next(), Some(0)); + assert_eq!(iter.next_back(), Some(0)); + assert_eq!(iter.advance_by(1), Ok(())); + assert_eq!(iter.advance_back_by(1), Ok(())); + assert_eq!(iter.len(), 96); +} + +#[test] +fn test_take_repeat_with_exact_size() { + let mut iter = repeat_with(|| 0).take(100); + assert_eq!(iter.len(), 100); + assert_eq!(iter.advance_by(100), Ok(())); + assert_eq!(iter.len(), 0); +}