Use iter::repeat_n to implement Vec::extend_with

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<T>`, short for 'vector'.
 ///
 /// # Examples
@@ -3258,31 +3264,7 @@ impl<T: Clone, A: Allocator> Vec<T, A> {
     #[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
-        }
+        <Self as SpecExtendWith<T>>::spec_extend_with(self, n, value);
     }
 }
 

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<T> {
+    #[track_caller]
+    fn spec_extend_with(&mut self, n: usize, value: T);
+}
+
+impl<T: Clone, A: Allocator> SpecExtendWith<T> for Vec<T, A> {
+    #[track_caller]
+    default fn spec_extend_with(&mut self, n: usize, value: T) {
+        self.extend_trusted(iter::repeat_n(value, n));
+    }
+}
+
+impl<T: Copy, A: Allocator> SpecExtendWith<T> for Vec<T, A> {
+    #[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) }
+    }
+}