Auto merge of rust-lang#120050 - scottmcm:vec-resize-memset, r=<try>

`Vec::resize` for bytes should be a single `memset`

Really I just started by trying to see if specializing `iter::repeat_n` would help the perf issue that kept me from removing `Vec::extend_with` last time I tried, but I noticed in the process that a resize for bytes doesn't set all the new space with a single `memset`: <https://play.rust-lang.org/?version=nightly&mode=release&edition=2021&gist=35175ec844b46fcd95e2d0aad526859e>

So using `repeat_n` to implement it -- like `VecDeque` uses, with the specialization for `next` to avoid a branch -- means that the optimizer notices the resize can set all the values with a single memset.

Author: bors

library/alloc/src/vec/mod.rs

@@ -2448,7 +2448,7 @@ impl<T: Clone, A: Allocator> Vec<T, A> {
         let len = self.len();
 
         if new_len > len {
-            self.extend_with(new_len - len, value)
+            self.extend_trusted(core::iter::repeat_n(value, new_len - len));
         } else {
             self.truncate(new_len);
         }
@@ -2562,38 +2562,6 @@ 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))]
-    /// 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
-        }
-    }
-}
-
 impl<T: PartialEq, A: Allocator> Vec<T, A> {
     /// Removes consecutive repeated elements in the vector according to the
     /// [`PartialEq`] trait implementation.
@@ -2919,7 +2887,7 @@ impl<T, A: Allocator> Vec<T, A> {
                     // 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);
+                    local_len.increment_len_unchecked(1);
                 });
             }
         } else {

library/alloc/src/vec/set_len_on_drop.rs

@@ -14,9 +14,12 @@ impl<'a> SetLenOnDrop<'a> {
         SetLenOnDrop { local_len: *len, len }
     }
 
+    /// # Safety
+    /// `self.current_len() + increment` must not overflow.
     #[inline]
-    pub(super) fn increment_len(&mut self, increment: usize) {
-        self.local_len += increment;
+    pub(super) unsafe fn increment_len_unchecked(&mut self, increment: usize) {
+        // SAFETY: This is our precondition
+        self.local_len = unsafe { self.local_len.unchecked_add(increment) };
     }
 
     #[inline]

library/alloc/src/vec/spec_from_elem.rs

@@ -13,7 +13,7 @@ pub(super) trait SpecFromElem: Sized {
 impl<T: Clone> SpecFromElem for T {
     default fn from_elem<A: Allocator>(elem: Self, n: usize, alloc: A) -> Vec<Self, A> {
         let mut v = Vec::with_capacity_in(n, alloc);
-        v.extend_with(n, elem);
+        v.extend_trusted(core::iter::repeat_n(elem, n));
         v
     }
 }
@@ -25,7 +25,7 @@ impl<T: Clone + IsZero> SpecFromElem for T {
             return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
         }
         let mut v = Vec::with_capacity_in(n, alloc);
-        v.extend_with(n, elem);
+        v.extend_trusted(core::iter::repeat_n(elem, n));
         v
     }
 }

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

@@ -110,6 +110,58 @@ impl<A> Drop for RepeatN<A> {
     }
 }
 
+trait SpecRepeatN<A> {
+    unsafe fn spec_next_unchecked(&mut self) -> A;
+    fn spec_fold<B, F: FnMut(B, A) -> B>(self, b: B, f: F) -> B;
+}
+
+impl<A: Clone> SpecRepeatN<A> for RepeatN<A> {
+    default unsafe fn spec_next_unchecked(&mut self) -> A {
+        self.count -= 1;
+
+        if self.count == 0 {
+            // SAFETY: 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(&self.element)
+        }
+    }
+
+    default fn spec_fold<B, F: FnMut(B, A) -> B>(mut self, mut b: B, mut f: F) -> B {
+        let mut count = self.count;
+        if let Some(element) = self.take_element() {
+            while count > 1 {
+                count -= 1;
+                b = f(b, element.clone());
+            }
+            b = f(b, element);
+        }
+        b
+    }
+}
+
+impl<A: Copy> SpecRepeatN<A> for RepeatN<A> {
+    unsafe fn spec_next_unchecked(&mut self) -> A {
+        self.count -= 1;
+
+        // For `Copy` types, we can always just read the item directly,
+        // so skip having a branch that would need to be optimized out.
+        *self.element
+    }
+
+    fn spec_fold<B, F: FnMut(B, A) -> B>(self, mut b: B, mut f: F) -> B {
+        let mut count = self.count;
+        let element = *self.element;
+
+        while count > 0 {
+            count -= 1;
+            b = f(b, element);
+        }
+        b
+    }
+}
+
 #[unstable(feature = "iter_repeat_n", issue = "104434")]
 impl<A: Clone> Iterator for RepeatN<A> {
     type Item = A;
@@ -120,15 +172,8 @@ impl<A: Clone> Iterator for RepeatN<A> {
             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(&self.element)
-        })
+        // SAFETY: Just confirmed above that the iterator is non-empty
+        unsafe { Some(self.spec_next_unchecked()) }
     }
 
     #[inline]
@@ -154,6 +199,11 @@ impl<A: Clone> Iterator for RepeatN<A> {
         }
     }
 
+    #[inline]
+    fn fold<B, F: FnMut(B, A) -> B>(self, init: B, f: F) -> B {
+        self.spec_fold(init, f)
+    }
+
     #[inline]
     fn last(mut self) -> Option<A> {
         self.take_element()

tests/codegen/vec-resize-memset.rs

@@ -0,0 +1,13 @@
+// compile-flags: -O
+// no-system-llvm
+// only-64bit
+// ignore-debug (the extra assertions get in the way)
+
+#![crate_type = "lib"]
+
+pub fn resize_with_bytes_is_one_memset(x: &mut Vec<u8>) {
+    let new_len = x.len() + 456789;
+    x.resize(new_len, 123);
+}
+
+// CHECK: call void @llvm.memset.p0.i64({{.+}}, i8 123, i64 456789, i1 false)