From d5c221107ec2bb1260f122f67c5217bc82ad90cc Mon Sep 17 00:00:00 2001 From: Waffle Date: Fri, 13 Nov 2020 12:50:12 +0300 Subject: [PATCH] add `Vec::extend_from_within` method Implement , 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()`). --- library/alloc/src/lib.rs | 1 + library/alloc/src/vec/mod.rs | 113 +++++++++++++++++++++++++++++++++-- library/alloc/tests/lib.rs | 1 + library/alloc/tests/vec.rs | 70 ++++++++++++++++++++++ 4 files changed, 181 insertions(+), 4 deletions(-) diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs index d0bfa038aa13d..2b4611d0fa0d1 100644 --- a/library/alloc/src/lib.rs +++ b/library/alloc/src/lib.rs @@ -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)] diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index ccc4f03a1e505..e71b9fe299095 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -1796,11 +1796,27 @@ impl Vec { #[unstable(feature = "vec_spare_capacity", issue = "75017")] #[inline] pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit] { + self.split_at_spare_mut().1 + } + + #[inline] + fn split_at_spare_mut(&mut self) -> (&mut [T], &mut [MaybeUninit]) { + 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` has the same layout as `T` + let spare_ptr = unsafe { ptr.cast::>().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, - 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) } } } @@ -1862,6 +1878,39 @@ impl Vec { 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(&mut self, src: R) + where + R: RangeBounds, + { + 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}`. @@ -1969,6 +2018,62 @@ pub fn from_elem_in(elem: T, n: usize, alloc: A) -> Vec< ::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); +} + +impl ExtendFromWithinSpec for Vec { + default unsafe fn spec_extend_from_within(&mut self, src: Range) { + 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 ExtendFromWithinSpec for Vec { + unsafe fn spec_extend_from_within(&mut self, src: Range) { + 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 //////////////////////////////////////////////////////////////////////////////// diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs index 0b7eeab4e9679..dd98f806451d8 100644 --- a/library/alloc/tests/lib.rs +++ b/library/alloc/tests/lib.rs @@ -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}; diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index e19406d7a0697..5c7ff67bc621c 100644 --- a/library/alloc/tests/vec.rs +++ b/library/alloc/tests/vec.rs @@ -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::::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"]); +}