diff --git a/library/alloc/src/collections/vec_deque/mod.rs b/library/alloc/src/collections/vec_deque/mod.rs index 85c809e0d188d..a988baf73ab26 100644 --- a/library/alloc/src/collections/vec_deque/mod.rs +++ b/library/alloc/src/collections/vec_deque/mod.rs @@ -2737,6 +2737,52 @@ impl Extend for VecDeque { self.push_back(elem); } + #[inline] + fn extend_from_array(&mut self, array: [A; N]) { + self.reserve(N); + + // Pre-leak the items in the array, since we'll duplicate them below. + // They'll never *actually* leak since nothing in the rest of the method can fail. + let array = ManuallyDrop::new(array); + + let available_at_head = self.cap() - self.head; + if available_at_head >= N { + // The easy case; we can copy it all at once. + + // SAFETY: We just checked we have enough space without overrunning capacity. + // Thus the additions are in-bounds and cannot wrap. + // The copy is non-overlapping since the input array is owned. + // Duplicating the items is fine as we've already kept them from being dropped. + unsafe { + let end = self.ptr().add(self.head); + ptr::copy_nonoverlapping(array.as_ptr(), end, N); + } + } else { + // The harder case; we need to copy some of the items to the end and some to the beginning. + + // SAFETY: We're copying exactly the number to go up to the capacity but not over. + // Thus the additions are in-bounds and cannot wrap. + // The copy is non-overlapping since the input array is owned. + // Duplicating the items is fine as we've already kept them from being dropped. + unsafe { + let end = self.ptr().add(self.head); + ptr::copy_nonoverlapping(array.as_ptr(), end, available_at_head); + } + + // SAFETY: The reserve ensured that we have space for N, and we're copying less than that. + // Thus the additions are in-bounds and cannot wrap. + // The copy is non-overlapping since the input array is owned. + // Duplicating the items is fine as we've already kept them from being dropped. + unsafe { + let remaining_to_copy = N - available_at_head; + let array_ptr = array.as_ptr().add(available_at_head); + ptr::copy_nonoverlapping(array_ptr, self.ptr(), remaining_to_copy); + } + } + + self.head = self.wrap_add(self.head, N); + } + #[inline] fn extend_reserve(&mut self, additional: usize) { self.reserve(additional); diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs index 3ac34c9ae28af..7c025aa97a757 100644 --- a/library/alloc/src/lib.rs +++ b/library/alloc/src/lib.rs @@ -99,6 +99,7 @@ #![feature(dropck_eyepatch)] #![feature(exact_size_is_empty)] #![feature(exclusive_range_pattern)] +#![feature(extend_from_array)] #![feature(extend_one)] #![feature(fmt_internals)] #![feature(fn_traits)] diff --git a/library/alloc/src/vec.rs b/library/alloc/src/vec.rs index 9fffb47aa5975..08a90eeb66753 100644 --- a/library/alloc/src/vec.rs +++ b/library/alloc/src/vec.rs @@ -2256,6 +2256,45 @@ impl Extend for Vec { self.push(item); } + /// Extends this vector by moving the elements of the array to the end of the vector. + /// + /// You can think of a call to + /// ```rust + /// #![feature(extend_from_array)] + /// # let mut v = vec![]; + /// v.extend_from_array([7, 13, 42]); + /// ``` + /// as though it were + /// ```rust + /// # let mut v = vec![]; + /// v.reserve(3); + /// v.push(7); + /// v.push(13); + /// v.push(42); + /// ``` + /// just simpler to type and easier on the optimizer. + /// + /// You should continue to use [`push`](Self::push) if you only have one element; + /// There's no advantage to doing `.extend_from_array([x])` instead. + #[inline] + fn extend_from_array(&mut self, array: [T; N]) { + self.reserve(N); + + // Pre-leak the items in the array, since we'll duplicate them below. + // They'll never *actually* leak since nothing in the rest of the method can fail. + let array = ManuallyDrop::new(array); + + // SAFETY: The reserve has guaranteed that we have enough space available. + // Thus the additions are in-bounds and cannot wrap. + // The copy is non-overlapping since the input array is owned. + // Duplicating the items is fine as we've already kept them from being dropped. + unsafe { + let end = self.as_mut_ptr().add(self.len); + ptr::copy_nonoverlapping(array.as_ptr(), end, N); + self.len += N; + } + } + #[inline] fn extend_reserve(&mut self, additional: usize) { self.reserve(additional); diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs index b7cc03f8eb999..2d788b5e4253a 100644 --- a/library/alloc/tests/lib.rs +++ b/library/alloc/tests/lib.rs @@ -4,6 +4,7 @@ #![feature(const_cow_is_borrowed)] #![feature(drain_filter)] #![feature(exact_size_is_empty)] +#![feature(extend_from_array)] #![feature(new_uninit)] #![feature(pattern)] #![feature(str_split_once)] diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index e19406d7a0697..acacc29f63d98 100644 --- a/library/alloc/tests/vec.rs +++ b/library/alloc/tests/vec.rs @@ -1954,3 +1954,17 @@ fn test_vec_swap() { assert_eq!(a[0], 42); assert_eq!(n, 0); } + +#[test] +fn test_vec_extend_from_array() { + let mut v = Vec::from("Hello"); + v.extend_from_array(*b" "); + v.extend_from_array(*b"World"); + v.push(b'!'); + assert_eq!(v, b"Hello World!"); + + // And use a Drop type to check for obvious use-after-free or similar + let mut v = vec![String::from("Hello")]; + v.extend_from_array([String::from(" "), String::from("world")]); + assert_eq!(v.concat(), "Hello world"); +} diff --git a/library/alloc/tests/vec_deque.rs b/library/alloc/tests/vec_deque.rs index 705f0d62fbb7a..4882047060d07 100644 --- a/library/alloc/tests/vec_deque.rs +++ b/library/alloc/tests/vec_deque.rs @@ -1728,3 +1728,34 @@ fn test_zero_sized_push() { } } } + +#[test] +fn test_vecdeque_extend_from_array() { + // Easy case where it's just adding to the end + let mut v = VecDeque::with_capacity(15); + v.extend_from_array(*b"12345"); + v.extend_from_array(*b"67890"); + v.extend_from_array(*b"abcde"); + assert_eq!(v, b"1234567890abcde"); + + // Check that the head hitting the end wraps it back to the correct place + let mut v = VecDeque::with_capacity(7); + v.extend_from_array(*b"1234"); + v.pop_front(); + v.extend_from_array(*b"5678"); + assert_eq!(v.as_slices(), (&b"2345678"[..], &b""[..])); + v.pop_front(); + v.push_back(b'!'); + assert_eq!(v.as_slices(), (&b"345678"[..], &b"!"[..])); + + // The version that needs two copies because it wraps around + let mut v = VecDeque::with_capacity(15); + v.extend_from_array(*b"1234567890"); + for _ in 0..5 { + v.pop_front().unwrap(); + } + v.extend_from_array(*b"ABCDEFGHIJ"); + assert_eq!(v, b"67890ABCDEFGHIJ"); + assert_eq!(v.capacity(), 15); + assert_eq!(v.as_slices(), (&b"67890ABCDEF"[..], &b"GHIJ"[..])); +} diff --git a/library/core/src/iter/traits/collect.rs b/library/core/src/iter/traits/collect.rs index 1ae6d15c12dd9..f7b26f7e497e9 100644 --- a/library/core/src/iter/traits/collect.rs +++ b/library/core/src/iter/traits/collect.rs @@ -341,6 +341,40 @@ pub trait Extend { self.extend(Some(item)); } + /// Extends a collection by moving the elements from the array. + /// + /// You should call this if you have multiple elements to add and have them all already, + /// as it allows the container to optimize that where possible. + /// + /// # Examples + /// + /// ```rust + /// #![feature(extend_from_array)] + /// let mut v = vec![0; 3]; + /// v.extend_from_array([1, 1, 2, 3, 5]); + /// assert_eq!(v, [0, 0, 0, 1, 1, 2, 3, 5]); + /// ``` + /// + /// Extending from an empty array is allowed, but of course doesn't affect the collection: + /// ```rust + /// #![feature(extend_from_array)] + /// let mut v = vec![3, 5, 4]; + /// v.extend_from_array([] as [i32; 0]); + /// assert_eq!(v, [3, 5, 4]); + /// ``` + /// + /// # Note to Implementors + /// + /// The default implementation of this will pass the array iterator to your `extend` implementation, + /// which is likely optimal for many cases. + /// + /// You should override this, however, if you can take advantage of the array elements being contiguous in memory. + /// (`Vec` does, for example, but there's no way to do so in `LinkedList`.) + #[unstable(feature = "extend_from_array", issue = "88888888")] + fn extend_from_array(&mut self, array: [A; N]) { + self.extend(crate::array::IntoIter::new(array)); + } + /// Reserves capacity in a collection for the given number of additional elements. /// /// The default implementation does nothing. diff --git a/src/test/codegen/vec-extend_from_array.rs b/src/test/codegen/vec-extend_from_array.rs new file mode 100644 index 0000000000000..64f2c8e63fd31 --- /dev/null +++ b/src/test/codegen/vec-extend_from_array.rs @@ -0,0 +1,25 @@ +// compile-flags: -O -C panic=abort +// ignore-debug: the debug assertions get in the way + +#![crate_type = "lib"] +#![feature(extend_from_array)] + +// This test is here to ensure that LLVM is optimizing out the `ManuallyDrop` construction, +// since we really don't want to copy the whole array to the stack and then again to the Vec. + +// CHECK-LABEL: @vec_extend_from_array_demo +#[no_mangle] +pub fn vec_extend_from_array_demo(v: &mut Vec, a: [String; 400]) { + // CHECK-NOT: alloca + // CHECK: call alloc::vec::Vec::reserve + // CHECK-NOT: alloca + // CHECK: call void @llvm.memcpy + // CHECK-NOT: alloca + // CHECK: ret + v.extend_from_array(a); +} + +// No validation against this one; it just keeps `reserve` from having only a single caller. +pub fn please_do_not_inline_the_reserve_call_llvm(v: &mut Vec, s: String) { + v.push(s); +}