diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 9856e9c18ec68..930fa74971299 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -150,6 +150,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_with::SpecExtendWith; + +#[cfg(not(no_global_oom_handling))] +mod spec_extend_with; + /// A contiguous growable array type, written as `Vec`, short for 'vector'. /// /// # Examples @@ -3258,31 +3264,7 @@ impl Vec { #[track_caller] /// 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 - } + >::spec_extend_with(self, n, value); } } diff --git a/library/alloc/src/vec/spec_extend_with.rs b/library/alloc/src/vec/spec_extend_with.rs new file mode 100644 index 0000000000000..6dd754e226512 --- /dev/null +++ b/library/alloc/src/vec/spec_extend_with.rs @@ -0,0 +1,36 @@ +use core::iter; +use core::mem::MaybeUninit; + +use super::Vec; +use crate::alloc::Allocator; + +// Specialization trait used for Vec::extend_with +pub(super) trait SpecExtendWith { + #[track_caller] + fn spec_extend_with(&mut self, n: usize, value: T); +} + +impl SpecExtendWith for Vec { + #[track_caller] + default fn spec_extend_with(&mut self, n: usize, value: T) { + self.extend_trusted(iter::repeat_n(value, n)); + } +} + +impl SpecExtendWith for Vec { + #[track_caller] + fn spec_extend_with(&mut self, n: usize, value: T) { + let len = self.len(); + self.reserve(n); + let unfilled = self.spare_capacity_mut(); + + // SAFETY: the above `reserve` call guarantees `n` to be in bounds. + let unfilled = unsafe { unfilled.get_unchecked_mut(..n) }; + for elem in unfilled { + *elem = MaybeUninit::new(value); + } + + // SAFETY: the elements have been initialized above. + unsafe { self.set_len(len + n) } + } +}