From 1697af255339f8a0e9751425009e51aefcd93b9b Mon Sep 17 00:00:00 2001 From: Scott McMurray Date: Tue, 16 Jan 2024 23:19:54 -0800 Subject: [PATCH] `Vec::resize` for bytes should be a single `memset` --- library/alloc/src/vec/mod.rs | 56 ++++++++++++++--------- library/alloc/src/vec/spec_extend_elem.rs | 27 +++++++++++ library/alloc/src/vec/spec_from_elem.rs | 19 +------- tests/codegen/vec-of-bytes-memset.rs | 23 ++++++++++ 4 files changed, 86 insertions(+), 39 deletions(-) create mode 100644 library/alloc/src/vec/spec_extend_elem.rs create mode 100644 tests/codegen/vec-of-bytes-memset.rs diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 2afb5dd0d1a26..b62e01fca6dc7 100644 --- a/library/alloc/src/vec/mod.rs +++ b/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`, short for 'vector'. /// /// # Examples @@ -2871,7 +2877,7 @@ impl Vec { 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 Vec<[T; N], A> { impl Vec { #[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 + 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; } } } diff --git a/library/alloc/src/vec/spec_extend_elem.rs b/library/alloc/src/vec/spec_extend_elem.rs new file mode 100644 index 0000000000000..851d6f8bbe33c --- /dev/null +++ b/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 { + fn spec_extend_elem(&mut self, n: usize, value: T); +} + +impl SpecExtendElem for Vec +where + T: Clone, +{ + default fn spec_extend_elem(&mut self, n: usize, value: T) { + self.extend_trusted(repeat_n(value, n)) + } +} + +impl SpecExtendElem for Vec +where + T: Copy, +{ + fn spec_extend_elem(&mut self, n: usize, value: T) { + self.extend_elem_copy(n, value) + } +} diff --git a/library/alloc/src/vec/spec_from_elem.rs b/library/alloc/src/vec/spec_from_elem.rs index 96d701e15d487..2b0e43fa8ceb7 100644 --- a/library/alloc/src/vec/spec_from_elem.rs +++ b/library/alloc/src/vec/spec_from_elem.rs @@ -12,7 +12,7 @@ pub(super) trait SpecFromElem: Sized { impl SpecFromElem for T { default fn from_elem(elem: Self, n: usize, alloc: A) -> Vec { 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 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(_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 - } -} diff --git a/tests/codegen/vec-of-bytes-memset.rs b/tests/codegen/vec-of-bytes-memset.rs new file mode 100644 index 0000000000000..2d5aedfd1b852 --- /dev/null +++ b/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) { + // 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 { + // 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] +}