Skip to content

Commit 50b394a

Browse files
scottmcmscottmcm
committed
Simplify Vec using iter::repeat_n
By using `repeat_n` instead of `ExtendElement`, we can delete a bunch more `unsafe` code from `Vec. All the codegen tests are passing (and I made one stricter to help be sure), so let's see if perf is also happy with it.
1 parent 71bb200 commit 50b394a

File tree

3 files changed

+37
-104
lines changed

3 files changed

+37
-104
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

+13-39
Original file line numberDiff line numberDiff line change
@@ -1,61 +1,35 @@
1-
use core::ptr;
1+
use core::iter;
22

33
use crate::alloc::Allocator;
44
use crate::raw_vec::RawVec;
55

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

88
// Specialization trait used for Vec::from_elem
99
pub(super) trait SpecFromElem: Sized {
1010
fn from_elem<A: Allocator>(elem: Self, n: usize, alloc: A) -> Vec<Self, A>;
1111
}
1212

13+
#[inline]
14+
fn basic_from_elem<T: Clone, A: Allocator>(elem: T, n: usize, alloc: A) -> Vec<T, A> {
15+
let mut v = Vec::with_capacity_in(n, alloc);
16+
v.extend(iter::repeat_n(elem, n));
17+
v
18+
}
19+
1320
impl<T: Clone> SpecFromElem for T {
21+
#[inline]
1422
default fn from_elem<A: Allocator>(elem: Self, n: usize, alloc: A) -> Vec<Self, A> {
15-
let mut v = Vec::with_capacity_in(n, alloc);
16-
v.extend_with(n, ExtendElement(elem));
17-
v
23+
basic_from_elem(elem, n, alloc)
1824
}
1925
}
2026

2127
impl<T: Clone + IsZero> SpecFromElem for T {
2228
#[inline]
23-
default fn from_elem<A: Allocator>(elem: T, n: usize, alloc: A) -> Vec<T, A> {
29+
fn from_elem<A: Allocator>(elem: T, n: usize, alloc: A) -> Vec<T, A> {
2430
if elem.is_zero() {
2531
return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
2632
}
27-
let mut v = Vec::with_capacity_in(n, alloc);
28-
v.extend_with(n, ExtendElement(elem));
29-
v
30-
}
31-
}
32-
33-
impl SpecFromElem for i8 {
34-
#[inline]
35-
fn from_elem<A: Allocator>(elem: i8, n: usize, alloc: A) -> Vec<i8, A> {
36-
if elem == 0 {
37-
return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
38-
}
39-
unsafe {
40-
let mut v = Vec::with_capacity_in(n, alloc);
41-
ptr::write_bytes(v.as_mut_ptr(), elem as u8, n);
42-
v.set_len(n);
43-
v
44-
}
45-
}
46-
}
47-
48-
impl SpecFromElem for u8 {
49-
#[inline]
50-
fn from_elem<A: Allocator>(elem: u8, n: usize, alloc: A) -> Vec<u8, A> {
51-
if elem == 0 {
52-
return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
53-
}
54-
unsafe {
55-
let mut v = Vec::with_capacity_in(n, alloc);
56-
ptr::write_bytes(v.as_mut_ptr(), elem, n);
57-
v.set_len(n);
58-
v
59-
}
33+
basic_from_elem(elem, n, alloc)
6034
}
6135
}

‎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)