Skip to content

Commit 4820629

Browse files
scottmcmscottmcm
committed
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 #104596, which tried to do too many things at once.)
1 parent 2ed65da commit 4820629

File tree

3 files changed

+28
-68
lines changed

3 files changed

+28
-68
lines changed

‎library/alloc/src/vec/mod.rs

+2-60
Original file line numberDiff line numberDiff line change
@@ -2163,7 +2163,7 @@ impl<T, A: Allocator> Vec<T, A> {
21632163
{
21642164
let len = self.len();
21652165
if new_len > len {
2166-
self.extend_with(new_len - len, ExtendFunc(f));
2166+
self.extend(iter::repeat_with(f).take(new_len - len));
21672167
} else {
21682168
self.truncate(new_len);
21692169
}
@@ -2361,7 +2361,7 @@ impl<T: Clone, A: Allocator> Vec<T, A> {
23612361
let len = self.len();
23622362

23632363
if new_len > len {
2364-
self.extend_with(new_len - len, ExtendElement(value))
2364+
self.extend(iter::repeat_n(value, new_len - len))
23652365
} else {
23662366
self.truncate(new_len);
23672367
}
@@ -2475,64 +2475,6 @@ impl<T, A: Allocator, const N: usize> Vec<[T; N], A> {
24752475
}
24762476
}
24772477

2478-
// This code generalizes `extend_with_{element,default}`.
2479-
trait ExtendWith<T> {
2480-
fn next(&mut self) -> T;
2481-
fn last(self) -> T;
2482-
}
2483-
2484-
struct ExtendElement<T>(T);
2485-
impl<T: Clone> ExtendWith<T> for ExtendElement<T> {
2486-
fn next(&mut self) -> T {
2487-
self.0.clone()
2488-
}
2489-
fn last(self) -> T {
2490-
self.0
2491-
}
2492-
}
2493-
2494-
struct ExtendFunc<F>(F);
2495-
impl<T, F: FnMut() -> T> ExtendWith<T> for ExtendFunc<F> {
2496-
fn next(&mut self) -> T {
2497-
(self.0)()
2498-
}
2499-
fn last(mut self) -> T {
2500-
(self.0)()
2501-
}
2502-
}
2503-
2504-
impl<T, A: Allocator> Vec<T, A> {
2505-
#[cfg(not(no_global_oom_handling))]
2506-
/// Extend the vector by `n` values, using the given generator.
2507-
fn extend_with<E: ExtendWith<T>>(&mut self, n: usize, mut value: E) {
2508-
self.reserve(n);
2509-
2510-
unsafe {
2511-
let mut ptr = self.as_mut_ptr().add(self.len());
2512-
// Use SetLenOnDrop to work around bug where compiler
2513-
// might not realize the store through `ptr` through self.set_len()
2514-
// don't alias.
2515-
let mut local_len = SetLenOnDrop::new(&mut self.len);
2516-
2517-
// Write all elements except the last one
2518-
for _ in 1..n {
2519-
ptr::write(ptr, value.next());
2520-
ptr = ptr.add(1);
2521-
// Increment the length in every step in case next() panics
2522-
local_len.increment_len(1);
2523-
}
2524-
2525-
if n > 0 {
2526-
// We can write the last element directly without cloning needlessly
2527-
ptr::write(ptr, value.last());
2528-
local_len.increment_len(1);
2529-
}
2530-
2531-
// len set by scope guard
2532-
}
2533-
}
2534-
}
2535-
25362478
impl<T: PartialEq, A: Allocator> Vec<T, A> {
25372479
/// Removes consecutive repeated elements in the vector according to the
25382480
/// [`PartialEq`] trait implementation.

‎library/alloc/src/vec/spec_from_elem.rs

+4-3
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1+
use core::iter;
12
use core::ptr;
23

34
use crate::alloc::Allocator;
45
use crate::raw_vec::RawVec;
56

6-
use super::{ExtendElement, IsZero, Vec};
7+
use super::{IsZero, Vec};
78

89
// Specialization trait used for Vec::from_elem
910
pub(super) trait SpecFromElem: Sized {
@@ -13,7 +14,7 @@ pub(super) trait SpecFromElem: Sized {
1314
impl<T: Clone> SpecFromElem for T {
1415
default fn from_elem<A: Allocator>(elem: Self, n: usize, alloc: A) -> Vec<Self, A> {
1516
let mut v = Vec::with_capacity_in(n, alloc);
16-
v.extend_with(n, ExtendElement(elem));
17+
v.extend(iter::repeat_n(elem, n));
1718
v
1819
}
1920
}
@@ -25,7 +26,7 @@ impl<T: Clone + IsZero> SpecFromElem for T {
2526
return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
2627
}
2728
let mut v = Vec::with_capacity_in(n, alloc);
28-
v.extend_with(n, ExtendElement(elem));
29+
v.extend(iter::repeat_n(elem, n));
2930
v
3031
}
3132
}

‎src/test/codegen/vec-calloc.rs

+22-5
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,15 @@
66
#![crate_type = "lib"]
77

88
// CHECK-LABEL: @vec_zero_bytes
9+
// CHECK-SAME: i64 %n
910
#[no_mangle]
1011
pub fn vec_zero_bytes(n: usize) -> Vec<u8> {
1112
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
1213
// CHECK-NOT: call {{.*}}reserve
1314
// CHECK-NOT: call {{.*}}__rust_alloc(
1415
// CHECK-NOT: call {{.*}}llvm.memset
1516

16-
// CHECK: call {{.*}}__rust_alloc_zeroed(
17+
// CHECK: call {{.*}}__rust_alloc_zeroed(i64 %n
1718

1819
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
1920
// CHECK-NOT: call {{.*}}reserve
@@ -22,17 +23,19 @@ pub fn vec_zero_bytes(n: usize) -> Vec<u8> {
2223

2324
// CHECK: ret void
2425
vec![0; n]
25-
}
2626

27+
}
2728
// CHECK-LABEL: @vec_one_bytes
29+
// CHECK-SAME: i64 %n
2830
#[no_mangle]
2931
pub fn vec_one_bytes(n: usize) -> Vec<u8> {
3032
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
3133
// CHECK-NOT: call {{.*}}reserve
3234
// CHECK-NOT: call {{.*}}__rust_alloc_zeroed(
3335

34-
// CHECK: call {{.*}}__rust_alloc(
36+
// CHECK: call {{.*}}__rust_alloc(i64 %n
3537
// CHECK: call {{.*}}llvm.memset
38+
// CHECK-SAME: i8 1, i64 %n
3639

3740
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
3841
// CHECK-NOT: call {{.*}}reserve
@@ -43,13 +46,20 @@ pub fn vec_one_bytes(n: usize) -> Vec<u8> {
4346
}
4447

4548
// CHECK-LABEL: @vec_zero_scalar
49+
// CHECK-SAME: i64 %n
4650
#[no_mangle]
4751
pub fn vec_zero_scalar(n: usize) -> Vec<i32> {
4852
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
4953
// CHECK-NOT: call {{.*}}reserve
5054
// CHECK-NOT: call {{.*}}__rust_alloc(
5155

52-
// CHECK: call {{.*}}__rust_alloc_zeroed(
56+
// CHECK: %[[BYTES:.+]] = shl i64 %n, 2
57+
58+
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
59+
// CHECK-NOT: call {{.*}}reserve
60+
// CHECK-NOT: call {{.*}}__rust_alloc(
61+
62+
// CHECK: call {{.*}}__rust_alloc_zeroed(i64 %[[BYTES]]
5363

5464
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
5565
// CHECK-NOT: call {{.*}}reserve
@@ -60,13 +70,20 @@ pub fn vec_zero_scalar(n: usize) -> Vec<i32> {
6070
}
6171

6272
// CHECK-LABEL: @vec_one_scalar
73+
// CHECK-SAME: i64 %n
6374
#[no_mangle]
6475
pub fn vec_one_scalar(n: usize) -> Vec<i32> {
6576
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
6677
// CHECK-NOT: call {{.*}}reserve
6778
// CHECK-NOT: call {{.*}}__rust_alloc_zeroed(
6879

69-
// CHECK: call {{.*}}__rust_alloc(
80+
// CHECK: %[[BYTES:.+]] = shl i64 %n, 2
81+
82+
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
83+
// CHECK-NOT: call {{.*}}reserve
84+
// CHECK-NOT: call {{.*}}__rust_alloc_zeroed(
85+
86+
// CHECK: call {{.*}}__rust_alloc(i64 %[[BYTES]]
7087

7188
// CHECK-NOT: call {{.*}}alloc::vec::from_elem
7289
// CHECK-NOT: call {{.*}}reserve

0 commit comments

Comments
 (0)