Use normal iterators instead of the internal vec::ExtendWith

Repeating is `TrustedLen`, so we don't need another copy of [the `TrustedLen` specialization for `extend`](https://github.com/rust-lang/rust/blob/c5d82ed7a4ad94a538bb87e5016e7d5ce0bd434b/library/alloc/src/vec/spec_extend.rs#L27). (This is a simpler version of rust-lang#104596, which tried to do too many things at once.)
scottmcm · Nov 20, 2022 · 4820629 · 4820629
1 parent 2ed65da
commit 4820629
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 68 deletions.
diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs
@@ -2163,7 +2163,7 @@ impl<T, A: Allocator> Vec<T, A> {
     {
         let len = self.len();
         if new_len > len {
-            self.extend_with(new_len - len, ExtendFunc(f));
+            self.extend(iter::repeat_with(f).take(new_len - len));
         } else {
             self.truncate(new_len);
         }
@@ -2361,7 +2361,7 @@ impl<T: Clone, A: Allocator> Vec<T, A> {
         let len = self.len();
 
         if new_len > len {
-            self.extend_with(new_len - len, ExtendElement(value))
+            self.extend(iter::repeat_n(value, new_len - len))
         } else {
             self.truncate(new_len);
         }
@@ -2475,64 +2475,6 @@ impl<T, A: Allocator, const N: usize> Vec<[T; N], A> {
     }
 }
 
-// This code generalizes `extend_with_{element,default}`.
-trait ExtendWith<T> {
-    fn next(&mut self) -> T;
-    fn last(self) -> T;
-}
-
-struct ExtendElement<T>(T);
-impl<T: Clone> ExtendWith<T> for ExtendElement<T> {
-    fn next(&mut self) -> T {
-        self.0.clone()
-    }
-    fn last(self) -> T {
-        self.0
-    }
-}
-
-struct ExtendFunc<F>(F);
-impl<T, F: FnMut() -> T> ExtendWith<T> for ExtendFunc<F> {
-    fn next(&mut self) -> T {
-        (self.0)()
-    }
-    fn last(mut self) -> T {
-        (self.0)()
-    }
-}
-
-impl<T, A: Allocator> Vec<T, A> {
-    #[cfg(not(no_global_oom_handling))]
-    /// Extend the vector by `n` values, using the given generator.
-    fn extend_with<E: ExtendWith<T>>(&mut self, n: usize, mut value: E) {
-        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.next());
-                ptr = ptr.add(1);
-                // Increment the length in every step in case next() panics
-                local_len.increment_len(1);
-            }
-
-            if n > 0 {
-                // We can write the last element directly without cloning needlessly
-                ptr::write(ptr, value.last());
-                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.

diff --git a/library/alloc/src/vec/spec_from_elem.rs b/library/alloc/src/vec/spec_from_elem.rs
@@ -1,9 +1,10 @@
+use core::iter;
 use core::ptr;
 
 use crate::alloc::Allocator;
 use crate::raw_vec::RawVec;
 
-use super::{ExtendElement, IsZero, Vec};
+use super::{IsZero, Vec};
 
 // Specialization trait used for Vec::from_elem
 pub(super) trait SpecFromElem: Sized {
@@ -13,7 +14,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, ExtendElement(elem));
+        v.extend(iter::repeat_n(elem, n));
         v
     }
 }
@@ -25,7 +26,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, ExtendElement(elem));
+        v.extend(iter::repeat_n(elem, n));
         v
     }
 }

diff --git a/src/test/codegen/vec-calloc.rs b/src/test/codegen/vec-calloc.rs
@@ -6,14 +6,15 @@
 #![crate_type = "lib"]
 
 // CHECK-LABEL: @vec_zero_bytes
+// CHECK-SAME: i64 %n
 #[no_mangle]
 pub fn vec_zero_bytes(n: usize) -> Vec<u8> {
     // CHECK-NOT: call {{.*}}alloc::vec::from_elem
     // CHECK-NOT: call {{.*}}reserve
     // CHECK-NOT: call {{.*}}__rust_alloc(
     // CHECK-NOT: call {{.*}}llvm.memset
 
-    // CHECK: call {{.*}}__rust_alloc_zeroed(
+    // CHECK: call {{.*}}__rust_alloc_zeroed(i64 %n
 
     // CHECK-NOT: call {{.*}}alloc::vec::from_elem
     // CHECK-NOT: call {{.*}}reserve
@@ -22,17 +23,19 @@ pub fn vec_zero_bytes(n: usize) -> Vec<u8> {
 
     // CHECK: ret void
     vec![0; n]
-}
 
+}
 // CHECK-LABEL: @vec_one_bytes
+// CHECK-SAME: i64 %n
 #[no_mangle]
 pub fn vec_one_bytes(n: usize) -> Vec<u8> {
     // CHECK-NOT: call {{.*}}alloc::vec::from_elem
     // CHECK-NOT: call {{.*}}reserve
     // CHECK-NOT: call {{.*}}__rust_alloc_zeroed(
 
-    // CHECK: call {{.*}}__rust_alloc(
+    // CHECK: call {{.*}}__rust_alloc(i64 %n
     // CHECK: call {{.*}}llvm.memset
+    // CHECK-SAME: i8 1, i64 %n
 
     // CHECK-NOT: call {{.*}}alloc::vec::from_elem
     // CHECK-NOT: call {{.*}}reserve
@@ -43,13 +46,20 @@ pub fn vec_one_bytes(n: usize) -> Vec<u8> {
 }
 
 // CHECK-LABEL: @vec_zero_scalar
+// CHECK-SAME: i64 %n
 #[no_mangle]
 pub fn vec_zero_scalar(n: usize) -> Vec<i32> {
     // CHECK-NOT: call {{.*}}alloc::vec::from_elem
     // CHECK-NOT: call {{.*}}reserve
     // CHECK-NOT: call {{.*}}__rust_alloc(
 
-    // CHECK: call {{.*}}__rust_alloc_zeroed(
+    // CHECK: %[[BYTES:.+]] = shl i64 %n, 2
+
+    // CHECK-NOT: call {{.*}}alloc::vec::from_elem
+    // CHECK-NOT: call {{.*}}reserve
+    // CHECK-NOT: call {{.*}}__rust_alloc(
+
+    // CHECK: call {{.*}}__rust_alloc_zeroed(i64 %[[BYTES]]
 
     // CHECK-NOT: call {{.*}}alloc::vec::from_elem
     // CHECK-NOT: call {{.*}}reserve
@@ -60,13 +70,20 @@ pub fn vec_zero_scalar(n: usize) -> Vec<i32> {
 }
 
 // CHECK-LABEL: @vec_one_scalar
+// CHECK-SAME: i64 %n
 #[no_mangle]
 pub fn vec_one_scalar(n: usize) -> Vec<i32> {
     // CHECK-NOT: call {{.*}}alloc::vec::from_elem
     // CHECK-NOT: call {{.*}}reserve
     // CHECK-NOT: call {{.*}}__rust_alloc_zeroed(
 
-    // CHECK: call {{.*}}__rust_alloc(
+    // CHECK: %[[BYTES:.+]] = shl i64 %n, 2
+
+    // CHECK-NOT: call {{.*}}alloc::vec::from_elem
+    // CHECK-NOT: call {{.*}}reserve
+    // CHECK-NOT: call {{.*}}__rust_alloc_zeroed(
+
+    // CHECK: call {{.*}}__rust_alloc(i64 %[[BYTES]]
 
     // CHECK-NOT: call {{.*}}alloc::vec::from_elem
     // CHECK-NOT: call {{.*}}reserve