Skip to content

Commit

Permalink
Simplify Vec using iter::repeat_n
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
scottmcm committed Nov 19, 2022
1 parent 71bb200 commit 50b394a
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 104 deletions.
62 changes: 2 additions & 60 deletions library/alloc/src/vec/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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.
Expand Down
52 changes: 13 additions & 39 deletions library/alloc/src/vec/spec_from_elem.rs
Original file line number Diff line number Diff line change
@@ -1,61 +1,35 @@
use core::ptr;
use core::iter;

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 {
fn from_elem<A: Allocator>(elem: Self, n: usize, alloc: A) -> Vec<Self, A>;
}

#[inline]
fn basic_from_elem<T: Clone, A: Allocator>(elem: T, n: usize, alloc: A) -> Vec<T, A> {
let mut v = Vec::with_capacity_in(n, alloc);
v.extend(iter::repeat_n(elem, n));
v
}

impl<T: Clone> SpecFromElem for T {
#[inline]
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
basic_from_elem(elem, n, alloc)
}
}

impl<T: Clone + IsZero> SpecFromElem for T {
#[inline]
default fn from_elem<A: Allocator>(elem: T, n: usize, alloc: A) -> Vec<T, A> {
fn from_elem<A: Allocator>(elem: T, n: usize, alloc: A) -> Vec<T, A> {
if elem.is_zero() {
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
}
}

impl SpecFromElem for i8 {
#[inline]
fn from_elem<A: Allocator>(elem: i8, n: usize, alloc: A) -> Vec<i8, A> {
if elem == 0 {
return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
}
unsafe {
let mut v = Vec::with_capacity_in(n, alloc);
ptr::write_bytes(v.as_mut_ptr(), elem as u8, n);
v.set_len(n);
v
}
}
}

impl SpecFromElem for u8 {
#[inline]
fn from_elem<A: Allocator>(elem: u8, n: usize, alloc: A) -> Vec<u8, A> {
if elem == 0 {
return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
}
unsafe {
let mut v = Vec::with_capacity_in(n, alloc);
ptr::write_bytes(v.as_mut_ptr(), elem, n);
v.set_len(n);
v
}
basic_from_elem(elem, n, alloc)
}
}
27 changes: 22 additions & 5 deletions src/test/codegen/vec-calloc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down

0 comments on commit 50b394a

Please sign in to comment.