Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Provides implementation of Vec::extend_from_slice optimized for T: Copy #236

Merged
merged 8 commits into from
Feb 21, 2024
71 changes: 71 additions & 0 deletions benches/benches.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,8 +104,78 @@ fn string_push_str(bump: &bumpalo::Bump, str: &str) {
criterion::black_box(s);
}

#[cfg(feature = "collections")]
fn extend_u8(bump: &bumpalo::Bump, slice: &[u8]) {
let slice = criterion::black_box(slice);
let mut vec = bumpalo::collections::Vec::<u8>::with_capacity_in(slice.len(), bump);
vec.extend(slice.iter().copied());
criterion::black_box(vec);
}

#[cfg(feature = "collections")]
fn extend_from_slice_u8(bump: &bumpalo::Bump, slice: &[u8]) {
let slice = criterion::black_box(slice);
let mut vec = bumpalo::collections::Vec::<u8>::with_capacity_in(slice.len(), bump);
vec.extend_from_slice(slice);
criterion::black_box(vec);
}

#[cfg(feature = "collections")]
fn extend_from_slice_copy_u8(bump: &bumpalo::Bump, slice: &[u8]) {
let slice = criterion::black_box(slice);
let mut vec = bumpalo::collections::Vec::<u8>::with_capacity_in(slice.len(), bump);
vec.extend_from_slice_copy(slice);
criterion::black_box(vec);
}

const ALLOCATIONS: usize = 10_000;

fn bench_extend_from_slice_copy(c: &mut Criterion) {
let lengths = &[
4usize,
5,
8,
11,
16,
64,
128,
331,
1024,
4 * 1024,
16 * 1024,
];

for len in lengths.iter().copied() {
let str = "x".repeat(len);
let mut group = c.benchmark_group(format!("extend {len} bytes"));
group.throughput(Throughput::Elements(len as u64));
group.bench_function("extend", |b| {
let mut bump = bumpalo::Bump::with_capacity(len);
b.iter(|| {
bump.reset();
extend_u8(&bump, str.as_bytes());
});
});
group.bench_function("extend_from_slice", |b| {
let mut bump = bumpalo::Bump::with_capacity(len);
let str = "x".repeat(len);
b.iter(|| {
bump.reset();
extend_from_slice_u8(&bump, str.as_bytes());
});
});
group.bench_function("extend_from_slice_copy", |b| {
let mut bump = bumpalo::Bump::with_capacity(len);
let str = "x".repeat(len);
b.iter(|| {
bump.reset();
extend_from_slice_copy_u8(&bump, str.as_bytes());
});
});
group.finish();
}
}

fn bench_alloc(c: &mut Criterion) {
let mut group = c.benchmark_group("alloc");
group.throughput(Throughput::Elements(ALLOCATIONS as u64));
Expand Down Expand Up @@ -249,6 +319,7 @@ fn bench_string_push_str(c: &mut Criterion) {

criterion_group!(
benches,
bench_extend_from_slice_copy,
bench_alloc,
bench_alloc_with,
bench_alloc_try_with,
Expand Down
26 changes: 1 addition & 25 deletions src/collections/string.rs
Original file line number Diff line number Diff line change
Expand Up @@ -936,31 +936,7 @@ impl<'bump> String<'bump> {
/// ```
#[inline]
pub fn push_str(&mut self, string: &str) {
// Reserve space in the Vec for the string to be added
let old_len = self.vec.len();
self.vec.reserve(string.len());

let new_len = old_len + string.len();
debug_assert!(new_len <= self.vec.capacity());

// Copy string into space just reserved
// SAFETY:
// * `src` is valid for reads of `string.len()` bytes by virtue of being an allocated `&str`.
// * `dst` is valid for writes of `string.len()` bytes as `self.vec.reserve(string.len())`
// above guarantees that.
// * Alignment is not relevant as `u8` has no alignment requirements.
// * Source and destination ranges cannot overlap as we just reserved the destination
// range from the bump.
unsafe {
let src = string.as_ptr();
let dst = self.vec.as_mut_ptr().add(old_len);
ptr::copy_nonoverlapping(src, dst, string.len());
}

// Update length of Vec to include string just pushed
// SAFETY: We reserved sufficent capacity for the string above.
// The elements at `old_len..new_len` were initialized by `copy_nonoverlapping` above.
unsafe { self.vec.set_len(new_len) };
self.vec.extend_from_slice_copy(string.as_bytes())
}

/// Returns this `String`'s capacity, in bytes.
Expand Down
65 changes: 63 additions & 2 deletions src/collections/vec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1777,6 +1777,67 @@ impl<'bump, T: 'bump + Clone> Vec<'bump, T> {
}
}

impl<'bump, T: 'bump + Copy> Vec<'bump, T> {
/// Copies all elements in the slice `other` and appends them to the `Vec`.
///
/// Note that this function is same as [`extend_from_slice`] except that it is optimized for
/// slices of types that implement the `Copy` trait. If and when Rust gets specialization
/// this function will likely be deprecated (but still available).
///
/// # Examples
///
/// ```
/// use bumpalo::{Bump, collections::Vec};
///
/// let b = Bump::new();
///
/// let mut vec = bumpalo::vec![in &b; 1];
/// vec.extend_from_slice_copy(&[2, 3, 4]);
/// assert_eq!(vec, [1, 2, 3, 4]);
/// ```
///
/// ```
/// use bumpalo::{Bump, collections::Vec};
///
/// let b = Bump::new();
///
/// let mut vec = bumpalo::vec![in &b; 'H' as u8];
/// vec.extend_from_slice_copy("ello, world!".as_bytes());
/// assert_eq!(vec, "Hello, world!".as_bytes());
/// ```
///
/// [`extend`]: #method.extend_from_slice
pub fn extend_from_slice_copy(&mut self, other: &[T]) {

// Reserve space in the Vec for the values to be added
let old_len = self.len();
self.reserve(other.len());

let new_len = old_len + other.len();
debug_assert!(new_len <= self.capacity());

// Copy values into the space that was just reserved
// SAFETY:
// * `src` is valid for reads of `other.len()` values by virtue of being a `&[T]`.
// * `dst` is valid for writes of `other.len()` bytes as `self.reserve(other.len())`
// above guarantees that.
// * Because `src` is a `&[T]` and dst is a `&[T]` within the `Vec<T>`,
// `copy_nonoverlapping`'s alignment requirements are met.
// * Source and destination ranges cannot overlap as we just reserved the destination
// range from the bump.
unsafe {
let src = other.as_ptr();
let dst = self.as_mut_ptr().add(old_len);
ptr::copy_nonoverlapping(src, dst, other.len());
}

// Update length of Vec to include values just pushed
// SAFETY: We reserved sufficient capacity for the values above.
// The elements at `old_len..new_len` were initialized by `copy_nonoverlapping` above.
unsafe { self.set_len(new_len) };
}
}

// This code generalises `extend_with_{element,default}`.
trait ExtendWith<T> {
fn next(&mut self) -> T;
Expand Down Expand Up @@ -2619,13 +2680,13 @@ where
impl<'bump> io::Write for Vec<'bump, u8> {
#[inline]
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.extend_from_slice(buf);
self.extend_from_slice_copy(buf);
Ok(buf.len())
}

#[inline]
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
self.extend_from_slice(buf);
self.extend_from_slice_copy(buf);
Ok(())
}

Expand Down