Skip to content

Commit

Permalink
add Vec::extend_from_within method
Browse files Browse the repository at this point in the history
Implement <rust-lang/rfcs#2714>, changes from the RFC:
- Rename the method `append_from_within` => `extend_from_within`
- Loose :Copy bound => :Clone
- Specialize in case of :Copy

This commit also adds `Vec::split_at_spare` private method and use it to implement
`Vec::spare_capacity_mut` and `Vec::extend_from_within`. This method returns 2
slices - initialized elements (same as `&mut vec[..]`) and uninitialized but
allocated space (same as `vec.spare_capacity_mut()`).
  • Loading branch information
WaffleLapkin committed Jan 31, 2021
1 parent 95cbcad commit d5c2211
Show file tree
Hide file tree
Showing 4 changed files with 181 additions and 4 deletions.
1 change: 1 addition & 0 deletions library/alloc/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@
#![cfg_attr(test, feature(test))]
#![cfg_attr(test, feature(new_uninit))]
#![feature(allocator_api)]
#![feature(vec_extend_from_within)]
#![feature(array_chunks)]
#![feature(array_methods)]
#![feature(array_value_iter)]
Expand Down
113 changes: 109 additions & 4 deletions library/alloc/src/vec/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1796,11 +1796,27 @@ impl<T, A: Allocator> Vec<T, A> {
#[unstable(feature = "vec_spare_capacity", issue = "75017")]
#[inline]
pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] {
self.split_at_spare_mut().1
}

#[inline]
fn split_at_spare_mut(&mut self) -> (&mut [T], &mut [MaybeUninit<T>]) {
let ptr = self.as_mut_ptr();

// Safety:
// - `ptr` is guaranteed to be in bounds for `capacity` elements
// - `len` is guaranteed to less or equal to `capacity`
// - `MaybeUninit<T>` has the same layout as `T`
let spare_ptr = unsafe { ptr.cast::<MaybeUninit<T>>().add(self.len) };

// Safety:
// - `ptr` is guaranteed to be valid for `len` elements
// - `spare_ptr` is offseted from `ptr` by `len`, so it doesn't overlap `initialized` slice
unsafe {
slice::from_raw_parts_mut(
self.as_mut_ptr().add(self.len) as *mut MaybeUninit<T>,
self.buf.capacity() - self.len,
)
let initialized = slice::from_raw_parts_mut(ptr, self.len);
let spare = slice::from_raw_parts_mut(spare_ptr, self.buf.capacity() - self.len);

(initialized, spare)
}
}
}
Expand Down Expand Up @@ -1862,6 +1878,39 @@ impl<T: Clone, A: Allocator> Vec<T, A> {
pub fn extend_from_slice(&mut self, other: &[T]) {
self.spec_extend(other.iter())
}

/// Copies elements from `src` range to the end of the vector.
///
/// ## Examples
///
/// ```
/// #![feature(vec_extend_from_within)]
///
/// let mut vec = vec![0, 1, 2, 3, 4];
///
/// vec.extend_from_within(2..);
/// assert_eq!(vec, [0, 1, 2, 3, 4, 2, 3, 4]);
///
/// vec.extend_from_within(..2);
/// assert_eq!(vec, [0, 1, 2, 3, 4, 2, 3, 4, 0, 1]);
///
/// vec.extend_from_within(4..8);
/// assert_eq!(vec, [0, 1, 2, 3, 4, 2, 3, 4, 0, 1, 4, 2, 3, 4]);
/// ```
#[unstable(feature = "vec_extend_from_within", issue = "none")]
pub fn extend_from_within<R>(&mut self, src: R)
where
R: RangeBounds<usize>,
{
let range = src.assert_len(self.len());
self.reserve(range.len());

// SAFETY:
// - `assert_len` guarantees that the given range is valid for indexing self
unsafe {
self.spec_extend_from_within(range);
}
}
}

// This code generalizes `extend_with_{element,default}`.
Expand Down Expand Up @@ -1969,6 +2018,62 @@ pub fn from_elem_in<T: Clone, A: Allocator>(elem: T, n: usize, alloc: A) -> Vec<
<T as SpecFromElem>::from_elem(elem, n, alloc)
}

trait ExtendFromWithinSpec {
/// Safety:
/// - `src` needs to be valid index
/// - `self.capacity() - self.len()` must be `>= src.len()`
unsafe fn spec_extend_from_within(&mut self, src: Range<usize>);
}

impl<T: Clone, A: Allocator> ExtendFromWithinSpec for Vec<T, A> {
default unsafe fn spec_extend_from_within(&mut self, src: Range<usize>) {
let initialized = {
let (this, spare) = self.split_at_spare_mut();

// Safety:
// - caller guaratees that src is a valid index
let to_clone = unsafe { this.get_unchecked(src) };

to_clone.iter().cloned().zip(spare.iter_mut()).map(|(e, s)| s.write(e)).count()
};

// Safety:
// - elements were just initialized
unsafe {
let new_len = self.len() + initialized;
self.set_len(new_len);
}
}
}

impl<T: Copy, A: Allocator> ExtendFromWithinSpec for Vec<T, A> {
unsafe fn spec_extend_from_within(&mut self, src: Range<usize>) {
let count = src.len();
{
let (init, spare) = self.split_at_spare_mut();

// Safety:
// - caller guaratees that `src` is a valid index
let source = unsafe { init.get_unchecked(src) };

// Safety:
// - Both pointers are created from unique slice references (`&mut [_]`)
// so they are valid and do not overlap.
// - Elements are :Copy so it's OK to to copy them, without doing
// anything with the original values
// - `count` is equal to the len of `source`, so source is valid for
// `count` reads
// - `.reserve(count)` guarantees that `spare.len() >= count` so spare
// is valid for `count` writes
unsafe { ptr::copy_nonoverlapping(source.as_ptr(), spare.as_mut_ptr() as _, count) };
}

// Safety:
// - The elements were just initialized by `copy_nonoverlapping`
self.len += count;
}
}

////////////////////////////////////////////////////////////////////////////////
// Common trait implementations for Vec
////////////////////////////////////////////////////////////////////////////////
Expand Down
1 change: 1 addition & 0 deletions library/alloc/tests/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
#![feature(int_bits_const)]
#![feature(vecdeque_binary_search)]
#![feature(slice_group_by)]
#![feature(vec_extend_from_within)]

use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
Expand Down
70 changes: 70 additions & 0 deletions library/alloc/tests/vec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1954,3 +1954,73 @@ fn test_vec_swap() {
assert_eq!(a[0], 42);
assert_eq!(n, 0);
}

#[test]
fn test_extend_from_within_spec() {
#[derive(Copy)]
struct CopyOnly;

impl Clone for CopyOnly {
fn clone(&self) -> Self {
panic!("extend_from_within must use specialization on copy");
}
}

vec![CopyOnly, CopyOnly].extend_from_within(..);
}

#[test]
fn test_extend_from_within_clone() {
let mut v = vec![String::from("sssss"), String::from("12334567890"), String::from("c")];
v.extend_from_within(1..);

assert_eq!(v, ["sssss", "12334567890", "c", "12334567890", "c"]);
}

#[test]
fn test_extend_from_within_complete_rande() {
let mut v = vec![0, 1, 2, 3];
v.extend_from_within(..);

assert_eq!(v, [0, 1, 2, 3, 0, 1, 2, 3]);
}

#[test]
fn test_extend_from_within_empty_rande() {
let mut v = vec![0, 1, 2, 3];
v.extend_from_within(1..1);

assert_eq!(v, [0, 1, 2, 3]);
}

#[test]
#[should_panic]
fn test_extend_from_within_out_of_rande() {
let mut v = vec![0, 1];
v.extend_from_within(..3);
}

#[test]
fn test_extend_from_within_zst() {
let mut v = vec![(); 8];
v.extend_from_within(3..7);

assert_eq!(v, [(); 12]);
}

#[test]
fn test_extend_from_within_empty_vec() {
let mut v = Vec::<i32>::new();
v.extend_from_within(..);

assert_eq!(v, []);
}

#[test]
fn test_extend_from_within() {
let mut v = vec![String::from("a"), String::from("b"), String::from("c")];
v.extend_from_within(1..=2);
v.extend_from_within(..=1);

assert_eq!(v, ["a", "b", "c", "b", "c", "a", "b"]);
}

0 comments on commit d5c2211

Please sign in to comment.