From 7e425b898540c51f73f152a353e5a138ae8602eb Mon Sep 17 00:00:00 2001 From: Antoni Spaanderman <56turtle56@gmail.com> Date: Fri, 2 Jan 2026 16:05:42 +0100 Subject: [PATCH] Add specialization for `deque1.prepend(deque2.drain(range))` This is used when moving elements between `VecDeque`s This pattern is also used in examples of `VecDeque::prepend`. (see its docs) --- .../alloc/src/collections/vec_deque/drain.rs | 6 +- .../src/collections/vec_deque/spec_extend.rs | 78 +++++++++++-- library/alloctests/tests/vec_deque.rs | 110 ++++++++++++++++++ 3 files changed, 180 insertions(+), 14 deletions(-) diff --git a/library/alloc/src/collections/vec_deque/drain.rs b/library/alloc/src/collections/vec_deque/drain.rs index eee7be674c953..a43853604a2d2 100644 --- a/library/alloc/src/collections/vec_deque/drain.rs +++ b/library/alloc/src/collections/vec_deque/drain.rs @@ -25,10 +25,10 @@ pub struct Drain< // drain_start is stored in deque.len pub(super) drain_len: usize, // index into the logical array, not the physical one (always lies in [0..deque.len)) - idx: usize, + pub(super) idx: usize, // number of elements after the drained range pub(super) tail_len: usize, - remaining: usize, + pub(super) remaining: usize, // Needed to make Drain covariant over T _marker: PhantomData<&'a T>, } @@ -53,7 +53,7 @@ impl<'a, T, A: Allocator> Drain<'a, T, A> { // Only returns pointers to the slices, as that's all we need // to drop them. May only be called if `self.remaining != 0`. - unsafe fn as_slices(&self) -> (*mut [T], *mut [T]) { + pub(super) unsafe fn as_slices(&self) -> (*mut [T], *mut [T]) { unsafe { let deque = self.deque.as_ref(); diff --git a/library/alloc/src/collections/vec_deque/spec_extend.rs b/library/alloc/src/collections/vec_deque/spec_extend.rs index f73ba795cbea4..4ef854e8b3275 100644 --- a/library/alloc/src/collections/vec_deque/spec_extend.rs +++ b/library/alloc/src/collections/vec_deque/spec_extend.rs @@ -1,7 +1,7 @@ use core::iter::{Copied, Rev, TrustedLen}; use core::slice; -use super::VecDeque; +use super::{Drain, VecDeque}; use crate::alloc::Allocator; #[cfg(not(test))] use crate::vec; @@ -157,7 +157,8 @@ impl SpecExtendFront> for VecDeque { #[track_caller] fn spec_extend_front(&mut self, mut iterator: vec::IntoIter) { let slice = iterator.as_slice(); - // SAFETY: elements in the slice are forgotten after this call + self.reserve(slice.len()); + // SAFETY: `slice.len()` space was just reserved and elements in the slice are forgotten after this call unsafe { prepend_reversed(self, slice) }; iterator.forget_remaining_elements(); } @@ -169,7 +170,8 @@ impl SpecExtendFront>> for VecDeque>) { let mut iterator = iterator.into_inner(); let slice = iterator.as_slice(); - // SAFETY: elements in the slice are forgotten after this call + self.reserve(slice.len()); + // SAFETY: `slice.len()` space was just reserved and elements in the slice are forgotten after this call unsafe { prepend(self, slice) }; iterator.forget_remaining_elements(); } @@ -182,7 +184,8 @@ where #[track_caller] fn spec_extend_front(&mut self, iter: Copied>) { let slice = iter.into_inner().as_slice(); - // SAFETY: T is Copy because Copied> is Iterator + self.reserve(slice.len()); + // SAFETY: `slice.len()` space was just reserved and T is Copy because Copied> is Iterator unsafe { prepend_reversed(self, slice) }; } } @@ -194,17 +197,69 @@ where #[track_caller] fn spec_extend_front(&mut self, iter: Rev>>) { let slice = iter.into_inner().into_inner().as_slice(); - // SAFETY: T is Copy because Rev>> is Iterator + self.reserve(slice.len()); + // SAFETY: `slice.len()` space was just reserved and T is Copy because Rev>> is Iterator unsafe { prepend(self, slice) }; } } +impl<'a, T, A1: Allocator, A2: Allocator> SpecExtendFront> for VecDeque { + #[track_caller] + fn spec_extend_front(&mut self, mut iter: Drain<'a, T, A2>) { + if iter.remaining == 0 { + return; + } + + self.reserve(iter.remaining); + unsafe { + // SAFETY: iter.remaining != 0. + let (left, right) = iter.as_slices(); + // SAFETY: + // - `iter.remaining` space was reserved, `iter.remaining == left.len() + right.len()`. + // - The elements in `left` and `right` are forgotten after these calls. + prepend_reversed(self, &*left); + prepend_reversed(self, &*right); + } + + iter.idx += iter.remaining; + iter.remaining = 0; + } +} + +impl<'a, T, A1: Allocator, A2: Allocator> SpecExtendFront>> + for VecDeque +{ + #[track_caller] + fn spec_extend_front(&mut self, iter: Rev>) { + let mut iter = iter.into_inner(); + + if iter.remaining == 0 { + return; + } + + self.reserve(iter.remaining); + unsafe { + // SAFETY: iter.remaining != 0. + let (left, right) = iter.as_slices(); + // SAFETY: + // - `iter.remaining` space was reserved, `iter.remaining == left.len() + right.len()`. + // - The elements in `left` and `right` are forgotten after these calls. + prepend(self, &*right); + prepend(self, &*left); + } + + iter.idx += iter.remaining; + iter.remaining = 0; + } +} + +/// Prepends elements of `slice` to `deque` using a copy. +/// /// # Safety /// -/// Elements of `slice` will be copied into the deque, make sure to forget the items if `T` is not `Copy`. +/// - `deque` must have space for `slice.len()` new elements. +/// - Elements of `slice` will be copied into the deque, make sure to forget the elements if `T` is not `Copy`. unsafe fn prepend(deque: &mut VecDeque, slice: &[T]) { - deque.reserve(slice.len()); - unsafe { deque.head = deque.wrap_sub(deque.head, slice.len()); deque.copy_slice(deque.head, slice); @@ -212,12 +267,13 @@ unsafe fn prepend(deque: &mut VecDeque, slice: &[T]) { } } +/// Prepends elements of `slice` to `deque` in reverse order using a copy. +/// /// # Safety /// -/// Elements of `slice` will be copied into the deque, make sure to forget the items if `T` is not `Copy`. +/// - `deque` must have space for `slice.len()` new elements. +/// - Elements of `slice` will be copied into the deque, make sure to forget the elements if `T` is not `Copy`. unsafe fn prepend_reversed(deque: &mut VecDeque, slice: &[T]) { - deque.reserve(slice.len()); - unsafe { deque.head = deque.wrap_sub(deque.head, slice.len()); deque.copy_slice_reversed(deque.head, slice); diff --git a/library/alloctests/tests/vec_deque.rs b/library/alloctests/tests/vec_deque.rs index 82803c7a0dabd..e9f860f8df9bd 100644 --- a/library/alloctests/tests/vec_deque.rs +++ b/library/alloctests/tests/vec_deque.rs @@ -2156,6 +2156,116 @@ fn test_extend_front_specialization_copy_slice() { assert_eq!(v.as_slices(), ([5].as_slice(), [4, 3, 2].as_slice())); } +#[test] +fn test_extend_front_specialization_deque_drain() { + // trigger 8 code paths: all combinations of prepend and extend_front, wrap and no wrap (src deque), wrap and no wrap (dst deque) + + /// Get deque containing `[1, 2, 3, 4]`, possibly wrapping in the middle (between the 2 and 3). + fn test_deque(wrap: bool) -> VecDeque { + if wrap { + let mut v = VecDeque::with_capacity(4); + v.extend([3, 4]); + v.prepend([1, 2]); + assert_eq!(v.as_slices(), ([1, 2].as_slice(), [3, 4].as_slice())); + v + } else { + VecDeque::from([1, 2, 3, 4]) + } + } + + // prepend, v2.head == 0 + + let mut v1 = VecDeque::with_capacity(7); + + let mut v2 = test_deque(false); + v1.prepend(v2.drain(..)); + // drain removes all elements but keeps the buffer + assert_eq!(v2, []); + assert!(v2.capacity() >= 4); + + assert_eq!(v1, [1, 2, 3, 4]); + v1.pop_back(); + + let mut v2 = test_deque(false); + // this should wrap around the physical buffer + v1.prepend(v2.drain(..)); + // drain removes all elements but keeps the buffer + assert_eq!(v2, []); + assert!(v2.capacity() >= 4); + + // check it really wrapped + assert_eq!(v1.as_slices(), ([1].as_slice(), [2, 3, 4, 1, 2, 3].as_slice())); + + // extend_front, v2.head == 0 + + let mut v1 = VecDeque::with_capacity(7); + + let mut v2 = test_deque(false); + v1.extend_front(v2.drain(..)); + // drain removes all elements but keeps the buffer + assert_eq!(v2, []); + assert!(v2.capacity() >= 4); + + assert_eq!(v1, [4, 3, 2, 1]); + v1.pop_back(); + + let mut v2 = test_deque(false); + // this should wrap around the physical buffer + v1.extend_front(v2.drain(..)); + // drain removes all elements but keeps the buffer + assert_eq!(v2, []); + assert!(v2.capacity() >= 4); + + // check it really wrapped + assert_eq!(v1.as_slices(), ([4].as_slice(), [3, 2, 1, 4, 3, 2].as_slice())); + + // prepend, v2.head != 0 + + let mut v1 = VecDeque::with_capacity(7); + + let mut v2 = test_deque(true); + v1.prepend(v2.drain(..)); + // drain removes all elements but keeps the buffer + assert_eq!(v2, []); + assert!(v2.capacity() >= 4); + + assert_eq!(v1, [1, 2, 3, 4]); + v1.pop_back(); + + let mut v2 = test_deque(true); + // this should wrap around the physical buffer + v1.prepend(v2.drain(..)); + // drain removes all elements but keeps the buffer + assert_eq!(v2, []); + assert!(v2.capacity() >= 4); + + // check it really wrapped + assert_eq!(v1.as_slices(), ([1].as_slice(), [2, 3, 4, 1, 2, 3].as_slice())); + + // extend_front, v2.head != 0 + + let mut v1 = VecDeque::with_capacity(7); + + let mut v2 = test_deque(true); + v1.extend_front(v2.drain(..)); + // drain removes all elements but keeps the buffer + assert_eq!(v2, []); + assert!(v2.capacity() >= 4); + + assert_eq!(v1, [4, 3, 2, 1]); + v1.pop_back(); + + let mut v2 = test_deque(true); + // this should wrap around the physical buffer + v1.extend_front(v2.drain(..)); + // drain removes all elements but keeps the buffer + assert_eq!(v2, []); + assert!(v2.capacity() >= 4); + + // check it really wrapped + assert_eq!(v1.as_slices(), ([4].as_slice(), [3, 2, 1, 4, 3, 2].as_slice())); +} + #[test] fn test_splice() { let mut v = VecDeque::from(vec![1, 2, 3, 4, 5]);