Auto merge of #149751 - scottmcm:spec-repeat-n-fold, r=<try>

Another attempt at `Vec::extend_with` via `iter::repeat_n`

Author: rust-bors[bot]

library/alloc/src/vec/mod.rs

@@ -3499,33 +3499,10 @@ impl<T, A: Allocator, const N: usize> Vec<[T; N], A> {
 
 impl<T: Clone, A: Allocator> Vec<T, A> {
     #[cfg(not(no_global_oom_handling))]
+    #[inline]
     /// Extend the vector by `n` clones of value.
     fn extend_with(&mut self, n: usize, value: T) {
-        self.reserve(n);
-
-        unsafe {
-            let mut ptr = self.as_mut_ptr().add(self.len());
-            // Use SetLenOnDrop to work around bug where compiler
-            // might not realize the store through `ptr` through self.set_len()
-            // don't alias.
-            let mut local_len = SetLenOnDrop::new(&mut self.len);
-
-            // Write all elements except the last one
-            for _ in 1..n {
-                ptr::write(ptr, value.clone());
-                ptr = ptr.add(1);
-                // Increment the length in every step in case clone() panics
-                local_len.increment_len(1);
-            }
-
-            if n > 0 {
-                // We can write the last element directly without cloning needlessly
-                ptr::write(ptr, value);
-                local_len.increment_len(1);
-            }
-
-            // len set by scope guard
-        }
+        self.extend_trusted(iter::repeat_n(value, n))
     }
 }
 
@@ -3907,16 +3884,29 @@ impl<T, A: Allocator> Vec<T, A> {
                 (low, high)
             );
             self.reserve(additional);
-            unsafe {
-                let ptr = self.as_mut_ptr();
-                let mut local_len = SetLenOnDrop::new(&mut self.len);
-                iterator.for_each(move |element| {
+
+            // Split out so it's generic only on T, not on the iterator (or allocator)
+            #[inline]
+            unsafe fn make_write_element_fn<T>(
+                ptr: *mut T,
+                mut local_len: SetLenOnDrop<'_>,
+            ) -> impl FnMut(T) {
+                move |element| unsafe {
                     ptr::write(ptr.add(local_len.current_len()), element);
                     // Since the loop executes user code which can panic we have to update
                     // the length every step to correctly drop what we've written.
                     // NB can't overflow since we would have had to alloc the address space
                     local_len.increment_len(1);
-                });
+                }
+            }
+
+            let ptr = self.as_mut_ptr();
+            let local_len = SetLenOnDrop::new(&mut self.len);
+
+            // SAFETY: we just allocated enough space and because it's `TrustedLen`
+            // we know it can't run extra times.
+            unsafe {
+                iterator.for_each(make_write_element_fn(ptr, local_len));
             }
         } else {
             // Per TrustedLen contract a `None` upper bound means that the iterator length

library/core/src/clone.rs

@@ -282,6 +282,11 @@ pub const trait Clone: Sized {
 #[marker]
 pub const unsafe trait TrivialClone: [const] Clone {}
 
+pub(crate) fn trivial_clone<T: TrivialClone>(x: &T) -> T {
+    // SAFETY: the point of the trait is that this is sound.
+    unsafe { crate::ptr::read(x) }
+}
+
 /// Derive macro generating an impl of the trait `Clone`.
 #[rustc_builtin_macro]
 #[stable(feature = "builtin_macro_prelude", since = "1.38.0")]

library/core/src/iter/sources/repeat_n.rs

@@ -1,3 +1,4 @@
+use crate::clone::{TrivialClone, trivial_clone};
 use crate::fmt;
 use crate::iter::{FusedIterator, TrustedLen, UncheckedIterator};
 use crate::num::NonZero;
@@ -153,6 +154,50 @@ impl<A: Clone> Iterator for RepeatN<A> {
     fn count(self) -> usize {
         self.len()
     }
+
+    #[inline]
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        if let Some(inner) = self.inner { SpecFold::spec_fold(inner, init, f) } else { init }
+    }
+}
+
+trait SpecFold: Sized {
+    fn spec_fold<B, F>(inner: RepeatNInner<Self>, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self) -> B;
+}
+
+impl<A: Clone> SpecFold for A {
+    default fn spec_fold<B, F>(inner: RepeatNInner<A>, mut init: B, mut f: F) -> B
+    where
+        F: FnMut(B, Self) -> B,
+    {
+        for _ in 1..inner.count.get() {
+            init = f(init, inner.element.clone());
+        }
+        f(init, inner.element)
+    }
+}
+
+impl<A: TrivialClone> SpecFold for A {
+    fn spec_fold<B, F>(inner: RepeatNInner<A>, mut init: B, mut f: F) -> B
+    where
+        F: FnMut(B, Self) -> B,
+    {
+        let RepeatNInner { element, count } = inner;
+        let mut n = count.get();
+        loop {
+            init = f(init, trivial_clone(&element));
+            n -= 1;
+            if n == 0 {
+                return init;
+            }
+        }
+    }
 }
 
 #[stable(feature = "iter_repeat_n", since = "1.82.0")]

library/coretests/tests/iter/sources.rs

@@ -79,6 +79,31 @@ fn test_repeat_with_take_collect() {
     assert_eq!(v, vec![1, 2, 4, 8, 16]);
 }
 
+#[test]
+fn test_repeat_n_trivial_clone_for_each() {
+    let mut n = 0;
+    repeat_n(123_u32, 10).for_each(|_| n += 1);
+    assert_eq!(n, 10);
+}
+
+#[test]
+fn test_repeat_n_custom_clone_for_each() {
+    #[derive(Copy)]
+    struct CloneCounter<'a>(&'a Cell<u32>);
+    impl Clone for CloneCounter<'_> {
+        fn clone(&self) -> Self {
+            self.0.set(self.0.get() + 1);
+            *self
+        }
+    }
+
+    let clone_count = Cell::new(0);
+    let mut n = 0;
+    repeat_n(CloneCounter(&clone_count), 10).for_each(|_| n += 1);
+    assert_eq!(n, 10);
+    assert_eq!(clone_count.get(), 9);
+}
+
 #[test]
 fn test_successors() {
     let mut powers_of_10 = successors(Some(1_u16), |n| n.checked_mul(10));