Merging upstream version 1.67.1+dfsg1.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

7 files changed, 311 insertions, 5 deletions

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<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        <Self as SpecFold>::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<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B;
+}
+
+impl<I, const N: usize> SpecFold for ArrayChunks<I, N>
+where
+    I: Iterator,
+{
+    #[inline]
+    default fn fold<B, F>(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<I, const N: usize> SpecFold for ArrayChunks<I, N>
+where
+    I: Iterator + TrustedRandomAccessNoCoerce,
+{
+    #[inline]
+    fn fold<B, F>(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<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
     where
-        Self: Sized,
         Fold: FnMut(Acc, Self::Item) -> R,
         R: Try<Output = Acc>,
     {
@@ -101,6 +100,26 @@ where
     impl_fold_via_try_fold! { fold -> try_fold }
 
     #[inline]
+    fn for_each<F: FnMut(Self::Item)>(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<usize> + '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> {
         let min = self.n.min(n);
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<i32> = 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<T: Clone>(element: T, count: usize) -> RepeatN<T> {
+    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<A> {
+    count: usize,
+    // Invariant: has been dropped iff count == 0.
+    element: ManuallyDrop<A>,
+}
+
+impl<A> RepeatN<A> {
+    /// 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<A> {
+        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<A> Drop for RepeatN<A> {
+    fn drop(&mut self) {
+        self.take_element();
+    }
+}
+
+#[unstable(feature = "iter_repeat_n", issue = "104434")]
+impl<A: Clone> Iterator for RepeatN<A> {
+    type Item = A;
+
+    #[inline]
+    fn next(&mut self) -> Option<A> {
+        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<usize>) {
+        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<A> {
+        self.take_element()
+    }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.len()
+    }
+}
+
+#[unstable(feature = "iter_repeat_n", issue = "104434")]
+impl<A: Clone> ExactSizeIterator for RepeatN<A> {
+    fn len(&self) -> usize {
+        self.count
+    }
+}
+
+#[unstable(feature = "iter_repeat_n", issue = "104434")]
+impl<A: Clone> DoubleEndedIterator for RepeatN<A> {
+    #[inline]
+    fn next_back(&mut self) -> Option<A> {
+        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<A> {
+        self.nth(n)
+    }
+}
+
+#[unstable(feature = "iter_repeat_n", issue = "104434")]
+impl<A: Clone> FusedIterator for RepeatN<A> {}
+
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<A: Clone> TrustedLen for RepeatN<A> {}
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, F: FnMut() -> A> Iterator for RepeatWith<F> {
     fn size_hint(&self) -> (usize, Option<usize>) {
         (usize::MAX, None)
     }
+
+    #[inline]
+    fn try_fold<Acc, Fold, R>(&mut self, mut init: Acc, mut fold: Fold) -> R
+    where
+        Fold: FnMut(Acc, Self::Item) -> R,
+        R: Try<Output = Acc>,
+    {
+        // 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<Item = ()>) {}
 
-/// 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