From 4547b622d8d29df964fa2914213088b148c498fc Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Wed, 17 Apr 2024 14:18:32 +0200 Subject: Merging upstream version 1.67.1+dfsg1. Signed-off-by: Daniel Baumann --- library/core/src/iter/adapters/array_chunks.rs | 75 +++++++++- library/core/src/iter/adapters/take.rs | 21 ++- library/core/src/iter/mod.rs | 2 + library/core/src/iter/sources.rs | 4 + library/core/src/iter/sources/repeat_n.rs | 195 +++++++++++++++++++++++++ library/core/src/iter/sources/repeat_with.rs | 17 +++ library/core/src/iter/traits/iterator.rs | 2 +- 7 files changed, 311 insertions(+), 5 deletions(-) create mode 100644 library/core/src/iter/sources/repeat_n.rs (limited to 'library/core/src/iter') diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index d4fb88610..5e4211058 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -1,6 +1,8 @@ use crate::array; -use crate::iter::{ByRefSized, FusedIterator, Iterator}; -use crate::ops::{ControlFlow, Try}; +use crate::const_closure::ConstFnMutClosure; +use crate::iter::{ByRefSized, FusedIterator, Iterator, TrustedRandomAccessNoCoerce}; +use crate::mem::{self, MaybeUninit}; +use crate::ops::{ControlFlow, NeverShortCircuit, Try}; /// An iterator over `N` elements of the iterator at a time. /// @@ -82,7 +84,13 @@ where } } - impl_fold_via_try_fold! { fold -> try_fold } + fn fold(self, init: B, f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + ::fold(self, init, f) + } } #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] @@ -168,3 +176,64 @@ where self.iter.len() < N } } + +trait SpecFold: Iterator { + fn fold(self, init: B, f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B; +} + +impl SpecFold for ArrayChunks +where + I: Iterator, +{ + #[inline] + default fn fold(mut self, init: B, mut f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + let fold = ConstFnMutClosure::new(&mut f, NeverShortCircuit::wrap_mut_2_imp); + self.try_fold(init, fold).0 + } +} + +impl SpecFold for ArrayChunks +where + I: Iterator + TrustedRandomAccessNoCoerce, +{ + #[inline] + fn fold(mut self, init: B, mut f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + let mut accum = init; + let inner_len = self.iter.size(); + let mut i = 0; + // Use a while loop because (0..len).step_by(N) doesn't optimize well. + while inner_len - i >= N { + let mut chunk = MaybeUninit::uninit_array(); + let mut guard = array::Guard { array_mut: &mut chunk, initialized: 0 }; + while guard.initialized < N { + // SAFETY: The method consumes the iterator and the loop condition ensures that + // all accesses are in bounds and only happen once. + unsafe { + let idx = i + guard.initialized; + guard.push_unchecked(self.iter.__iterator_get_unchecked(idx)); + } + } + mem::forget(guard); + // SAFETY: The loop above initialized all elements + let chunk = unsafe { MaybeUninit::array_assume_init(chunk) }; + accum = f(accum, chunk); + i += N; + } + + // unlike try_fold this method does not need to take care of the remainder + // since `self` will be dropped + + accum + } +} diff --git a/library/core/src/iter/adapters/take.rs b/library/core/src/iter/adapters/take.rs index 58a0b9d7b..d947c7b0e 100644 --- a/library/core/src/iter/adapters/take.rs +++ b/library/core/src/iter/adapters/take.rs @@ -75,7 +75,6 @@ where #[inline] fn try_fold(&mut self, init: Acc, fold: Fold) -> R where - Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try, { @@ -100,6 +99,26 @@ where impl_fold_via_try_fold! { fold -> try_fold } + #[inline] + fn for_each(mut self, f: F) { + // The default implementation would use a unit accumulator, so we can + // avoid a stateful closure by folding over the remaining number + // of items we wish to return instead. + fn check<'a, Item>( + mut action: impl FnMut(Item) + 'a, + ) -> impl FnMut(usize, Item) -> Option + 'a { + move |more, x| { + action(x); + more.checked_sub(1) + } + } + + let remaining = self.n; + if remaining > 0 { + self.iter.try_fold(remaining - 1, check(f)); + } + } + #[inline] #[rustc_inherit_overflow_checks] fn advance_by(&mut self, n: usize) -> Result<(), usize> { diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs index ef0f39782..bb35d50b4 100644 --- a/library/core/src/iter/mod.rs +++ b/library/core/src/iter/mod.rs @@ -401,6 +401,8 @@ pub use self::sources::{once, Once}; pub use self::sources::{once_with, OnceWith}; #[stable(feature = "rust1", since = "1.0.0")] pub use self::sources::{repeat, Repeat}; +#[unstable(feature = "iter_repeat_n", issue = "104434")] +pub use self::sources::{repeat_n, RepeatN}; #[stable(feature = "iterator_repeat_with", since = "1.28.0")] pub use self::sources::{repeat_with, RepeatWith}; #[stable(feature = "iter_successors", since = "1.34.0")] diff --git a/library/core/src/iter/sources.rs b/library/core/src/iter/sources.rs index d34772cd3..3ec426a3a 100644 --- a/library/core/src/iter/sources.rs +++ b/library/core/src/iter/sources.rs @@ -4,6 +4,7 @@ mod from_generator; mod once; mod once_with; mod repeat; +mod repeat_n; mod repeat_with; mod successors; @@ -16,6 +17,9 @@ pub use self::empty::{empty, Empty}; #[stable(feature = "iter_once", since = "1.2.0")] pub use self::once::{once, Once}; +#[unstable(feature = "iter_repeat_n", issue = "104434")] +pub use self::repeat_n::{repeat_n, RepeatN}; + #[stable(feature = "iterator_repeat_with", since = "1.28.0")] pub use self::repeat_with::{repeat_with, RepeatWith}; diff --git a/library/core/src/iter/sources/repeat_n.rs b/library/core/src/iter/sources/repeat_n.rs new file mode 100644 index 000000000..fd8d25ce1 --- /dev/null +++ b/library/core/src/iter/sources/repeat_n.rs @@ -0,0 +1,195 @@ +use crate::iter::{FusedIterator, TrustedLen}; +use crate::mem::ManuallyDrop; + +/// Creates a new iterator that repeats a single element a given number of times. +/// +/// The `repeat_n()` function repeats a single value exactly `n` times. +/// +/// This is very similar to using [`repeat()`] with [`Iterator::take()`], +/// but there are two differences: +/// - `repeat_n()` can return the original value, rather than always cloning. +/// - `repeat_n()` produces an [`ExactSizeIterator`]. +/// +/// [`repeat()`]: crate::iter::repeat +/// +/// # Examples +/// +/// Basic usage: +/// +/// ``` +/// #![feature(iter_repeat_n)] +/// use std::iter; +/// +/// // four of the number four: +/// let mut four_fours = iter::repeat_n(4, 4); +/// +/// assert_eq!(Some(4), four_fours.next()); +/// assert_eq!(Some(4), four_fours.next()); +/// assert_eq!(Some(4), four_fours.next()); +/// assert_eq!(Some(4), four_fours.next()); +/// +/// // no more fours +/// assert_eq!(None, four_fours.next()); +/// ``` +/// +/// For non-`Copy` types, +/// +/// ``` +/// #![feature(iter_repeat_n)] +/// use std::iter; +/// +/// let v: Vec = Vec::with_capacity(123); +/// let mut it = iter::repeat_n(v, 5); +/// +/// for i in 0..4 { +/// // It starts by cloning things +/// let cloned = it.next().unwrap(); +/// assert_eq!(cloned.len(), 0); +/// assert_eq!(cloned.capacity(), 0); +/// } +/// +/// // ... but the last item is the original one +/// let last = it.next().unwrap(); +/// assert_eq!(last.len(), 0); +/// assert_eq!(last.capacity(), 123); +/// +/// // ... and now we're done +/// assert_eq!(None, it.next()); +/// ``` +#[inline] +#[unstable(feature = "iter_repeat_n", issue = "104434")] +#[doc(hidden)] // waiting on ACP#120 to decide whether to expose publicly +pub fn repeat_n(element: T, count: usize) -> RepeatN { + let mut element = ManuallyDrop::new(element); + + if count == 0 { + // SAFETY: we definitely haven't dropped it yet, since we only just got + // passed it in, and because the count is zero the instance we're about + // to create won't drop it, so to avoid leaking we need to now. + unsafe { ManuallyDrop::drop(&mut element) }; + } + + RepeatN { element, count } +} + +/// An iterator that repeats an element an exact number of times. +/// +/// This `struct` is created by the [`repeat_n()`] function. +/// See its documentation for more. +#[derive(Clone, Debug)] +#[unstable(feature = "iter_repeat_n", issue = "104434")] +#[doc(hidden)] // waiting on ACP#120 to decide whether to expose publicly +pub struct RepeatN { + count: usize, + // Invariant: has been dropped iff count == 0. + element: ManuallyDrop, +} + +impl RepeatN { + /// If we haven't already dropped the element, return it in an option. + /// + /// Clears the count so it won't be dropped again later. + #[inline] + fn take_element(&mut self) -> Option { + if self.count > 0 { + self.count = 0; + // SAFETY: We just set count to zero so it won't be dropped again, + // and it used to be non-zero so it hasn't already been dropped. + unsafe { Some(ManuallyDrop::take(&mut self.element)) } + } else { + None + } + } +} + +#[unstable(feature = "iter_repeat_n", issue = "104434")] +impl Drop for RepeatN { + fn drop(&mut self) { + self.take_element(); + } +} + +#[unstable(feature = "iter_repeat_n", issue = "104434")] +impl Iterator for RepeatN { + type Item = A; + + #[inline] + fn next(&mut self) -> Option { + if self.count == 0 { + return None; + } + + self.count -= 1; + Some(if self.count == 0 { + // SAFETY: the check above ensured that the count used to be non-zero, + // so element hasn't been dropped yet, and we just lowered the count to + // zero so it won't be dropped later, and thus it's okay to take it here. + unsafe { ManuallyDrop::take(&mut self.element) } + } else { + A::clone(&mut self.element) + }) + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let len = self.len(); + (len, Some(len)) + } + + #[inline] + fn advance_by(&mut self, skip: usize) -> Result<(), usize> { + let len = self.count; + + if skip >= len { + self.take_element(); + } + + if skip > len { + Err(len) + } else { + self.count = len - skip; + Ok(()) + } + } + + #[inline] + fn last(mut self) -> Option { + self.take_element() + } + + #[inline] + fn count(self) -> usize { + self.len() + } +} + +#[unstable(feature = "iter_repeat_n", issue = "104434")] +impl ExactSizeIterator for RepeatN { + fn len(&self) -> usize { + self.count + } +} + +#[unstable(feature = "iter_repeat_n", issue = "104434")] +impl DoubleEndedIterator for RepeatN { + #[inline] + fn next_back(&mut self) -> Option { + self.next() + } + + #[inline] + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + self.advance_by(n) + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option { + self.nth(n) + } +} + +#[unstable(feature = "iter_repeat_n", issue = "104434")] +impl FusedIterator for RepeatN {} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl TrustedLen for RepeatN {} diff --git a/library/core/src/iter/sources/repeat_with.rs b/library/core/src/iter/sources/repeat_with.rs index 6f62662d8..ab2d0472b 100644 --- a/library/core/src/iter/sources/repeat_with.rs +++ b/library/core/src/iter/sources/repeat_with.rs @@ -1,4 +1,5 @@ use crate::iter::{FusedIterator, TrustedLen}; +use crate::ops::Try; /// Creates a new iterator that repeats elements of type `A` endlessly by /// applying the provided closure, the repeater, `F: FnMut() -> A`. @@ -89,6 +90,22 @@ impl A> Iterator for RepeatWith { fn size_hint(&self) -> (usize, Option) { (usize::MAX, None) } + + #[inline] + fn try_fold(&mut self, mut init: Acc, mut fold: Fold) -> R + where + Fold: FnMut(Acc, Self::Item) -> R, + R: Try, + { + // This override isn't strictly needed, but avoids the need to optimize + // away the `next`-always-returns-`Some` and emphasizes that the `?` + // is the only way to exit the loop. + + loop { + let item = (self.repeater)(); + init = fold(init, item)?; + } + } } #[stable(feature = "iterator_repeat_with", since = "1.28.0")] diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs index 789a87968..83c7e8977 100644 --- a/library/core/src/iter/traits/iterator.rs +++ b/library/core/src/iter/traits/iterator.rs @@ -14,7 +14,7 @@ use super::super::{ fn _assert_is_object_safe(_: &dyn Iterator) {} -/// An interface for dealing with iterators. +/// A trait for dealing with iterators. /// /// This is the main iterator trait. For more about the concept of iterators /// generally, please see the [module-level documentation]. In particular, you -- cgit v1.2.3