Vec::resize for bytes should be a single memset

library/alloc/src/vec/mod.rs

@@ -145,6 +145,12 @@ use self::spec_extend::SpecExtend;
 #[cfg(not(no_global_oom_handling))]
 mod spec_extend;
 
+#[cfg(not(no_global_oom_handling))]
+use self::spec_extend_elem::SpecExtendElem;
+
+#[cfg(not(no_global_oom_handling))]
+mod spec_extend_elem;
+
 /// A contiguous growable array type, written as `Vec<T>`, short for 'vector'.
 ///
 /// # Examples
@@ -2871,7 +2877,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_elem(new_len - len, value)
  } else {
  self.truncate(new_len);
  }
@@ -2985,32 +2991,38 @@ 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) {
+ #[inline]
+ /// Extend the vector by `n` clones of `value`.
+ ///
+ /// Uses specialization to dispatch to `extend_elem_copy` if possible.
+ fn extend_elem(&mut self, n: usize, value: T) {
+ self.spec_extend_elem(n, value);
+ }
+
+ #[cfg(not(no_global_oom_handling))]
+ #[inline]
+ /// Extend the vector by `n` *copies* of `value`.
+ fn extend_elem_copy(&mut self, n: usize, value: T)
+ where
+ T: Copy,
+ {
  self.reserve(n);
 
+ // SAFETY: We now have space for all the elements, so the pointer math
+ // is all in-bounds. Because `T: Copy`, there's no user code being run
+ // here (no `clone` nor any `drop` in the assignment).
  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);
+ // Because there's no user code being run here, we can skip it for ZSTs.
+ // That helps tests in debug mode that do things like `vec![(); HUGE]`.
+ // See <https:/rust-lang/rust/pull/118094>
+ if !T::IS_ZST {
+ let ptr = self.as_mut_ptr().add(self.len);
+ for i in 0..n {
+ *ptr.add(i) = value;
+ }
  }
 
- // len set by scope guard
+ self.len += n;
  }
  }
 }

library/alloc/src/vec/spec_extend_elem.rs

@@ -0,0 +1,27 @@
+use core::iter::repeat_n;
+
+use super::Vec;
+use crate::alloc::Allocator;
+
+// Specialization trait used for Vec::extend_elem
+pub(super) trait SpecExtendElem<T> {
+ fn spec_extend_elem(&mut self, n: usize, value: T);
+}
+
+impl<T, A: Allocator> SpecExtendElem<T> for Vec<T, A>
+where
+ T: Clone,
+{
+ default fn spec_extend_elem(&mut self, n: usize, value: T) {
+ self.extend_trusted(repeat_n(value, n))
+ }
+}
+
+impl<T, A: Allocator> SpecExtendElem<T> for Vec<T, A>
+where
+ T: Copy,
+{
+ fn spec_extend_elem(&mut self, n: usize, value: T) {
+ self.extend_elem_copy(n, value)
+ }
+}

library/alloc/src/vec/spec_from_elem.rs

@@ -12,7 +12,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_elem(n, elem);
  v
  }
 }
@@ -24,7 +24,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_elem(n, elem);
  v
  }
 }
@@ -58,18 +58,3 @@ impl SpecFromElem for u8 {
  v
  }
 }
-
-// A better way would be to implement this for all ZSTs which are `Copy` and have trivial `Clone`
-// but the latter cannot be detected currently
-impl SpecFromElem for () {
- #[inline]
- fn from_elem<A: Allocator>(_elem: (), n: usize, alloc: A) -> Vec<(), A> {
- let mut v = Vec::with_capacity_in(n, alloc);
- // SAFETY: the capacity has just been set to `n`
- // and `()` is a ZST with trivial `Clone` implementation
- unsafe {
- v.set_len(n);
- }
- v
- }
-}

tests/codegen/vec-of-bytes-memset.rs

@@ -0,0 +1,23 @@
+//@ compile-flags: -O
+//@ only-64bit
+
+#![crate_type = "lib"]
+
+// CHECK-LABEL: @resize_bytes_is_one_memset
+#[no_mangle]
+pub fn resize_bytes_is_one_memset(x: &mut Vec<u8>) {
+ // CHECK: call void @llvm.memset.p0.i64({{.+}}, i8 123, i64 456789, i1 false)
+ let new_len = x.len() + 456789;
+ x.resize(new_len, 123);
+}
+
+#[derive(Copy, Clone)]
+struct ByteNewtype(i8);
+
+// CHECK-LABEL: @from_elem_is_one_memset
+#[no_mangle]
+pub fn from_elem_is_one_memset() -> Vec<ByteNewtype> {
+ // CHECK: %[[P:.+]] = tail call{{.+}}@__rust_alloc(i64 noundef 123456, i64 noundef 1)
+ // CHECK: call void @llvm.memset.p0.i64({{.+}} %[[P]], i8 42, i64 123456, i1 false)
+ vec![ByteNewtype(42); 123456]
+}