From a9f20e315a1c6334f6e93db58366a415d02ae0ec Mon Sep 17 00:00:00 2001 From: Jeremy Smart Date: Wed, 1 Oct 2025 18:29:30 -0400 Subject: [PATCH] add VecDeque::splice --- .../alloc/src/collections/vec_deque/drain.rs | 23 ++- .../alloc/src/collections/vec_deque/mod.rs | 65 ++++++++ .../alloc/src/collections/vec_deque/splice.rs | 143 ++++++++++++++++++ library/alloctests/tests/lib.rs | 1 + library/alloctests/tests/vec_deque.rs | 71 +++++++++ 5 files changed, 291 insertions(+), 12 deletions(-) create mode 100644 library/alloc/src/collections/vec_deque/splice.rs diff --git a/library/alloc/src/collections/vec_deque/drain.rs b/library/alloc/src/collections/vec_deque/drain.rs index 321621d18be9f..eee7be674c953 100644 --- a/library/alloc/src/collections/vec_deque/drain.rs +++ b/library/alloc/src/collections/vec_deque/drain.rs @@ -21,13 +21,13 @@ pub struct Drain< > { // We can't just use a &mut VecDeque, as that would make Drain invariant over T // and we want it to be covariant instead - deque: NonNull>, + pub(super) deque: NonNull>, // drain_start is stored in deque.len - drain_len: usize, + pub(super) drain_len: usize, // index into the logical array, not the physical one (always lies in [0..deque.len)) idx: usize, - // number of elements remaining after dropping the drain - new_len: usize, + // number of elements after the drained range + pub(super) tail_len: usize, remaining: usize, // Needed to make Drain covariant over T _marker: PhantomData<&'a T>, @@ -40,12 +40,12 @@ impl<'a, T, A: Allocator> Drain<'a, T, A> { drain_len: usize, ) -> Self { let orig_len = mem::replace(&mut deque.len, drain_start); - let new_len = orig_len - drain_len; + let tail_len = orig_len - drain_start - drain_len; Drain { deque: NonNull::from(deque), drain_len, idx: drain_start, - new_len, + tail_len, remaining: drain_len, _marker: PhantomData, } @@ -78,7 +78,7 @@ impl fmt::Debug for Drain<'_, T, A> { f.debug_tuple("Drain") .field(&self.drain_len) .field(&self.idx) - .field(&self.new_len) + .field(&self.tail_len) .field(&self.remaining) .finish() } @@ -125,7 +125,9 @@ impl Drop for Drain<'_, T, A> { let source_deque = unsafe { self.0.deque.as_mut() }; let drain_len = self.0.drain_len; - let new_len = self.0.new_len; + let head_len = source_deque.len; // #elements in front of the drain + let tail_len = self.0.tail_len; // #elements behind the drain + let new_len = head_len + tail_len; if T::IS_ZST { // no need to copy around any memory if T is a ZST @@ -133,9 +135,6 @@ impl Drop for Drain<'_, T, A> { return; } - let head_len = source_deque.len; // #elements in front of the drain - let tail_len = new_len - head_len; // #elements behind the drain - // Next, we will fill the hole left by the drain with as few writes as possible. // The code below handles the following control flow and reduces the amount of // branches under the assumption that `head_len == 0 || tail_len == 0`, i.e. @@ -219,7 +218,7 @@ impl Drop for Drain<'_, T, A> { } if new_len == 0 { - // Special case: If the entire dequeue was drained, reset the head back to 0, + // Special case: If the entire deque was drained, reset the head back to 0, // like `.clear()` does. source_deque.head = 0; } else if head_len < tail_len { diff --git a/library/alloc/src/collections/vec_deque/mod.rs b/library/alloc/src/collections/vec_deque/mod.rs index ac619a42d356d..814cf958b337d 100644 --- a/library/alloc/src/collections/vec_deque/mod.rs +++ b/library/alloc/src/collections/vec_deque/mod.rs @@ -55,6 +55,13 @@ use self::spec_from_iter::SpecFromIter; mod spec_from_iter; +#[cfg(not(no_global_oom_handling))] +#[unstable(feature = "deque_extend_front", issue = "146975")] +pub use self::splice::Splice; + +#[cfg(not(no_global_oom_handling))] +mod splice; + #[cfg(test)] mod tests; @@ -1625,6 +1632,64 @@ impl VecDeque { unsafe { Drain::new(self, drain_start, drain_len) } } + /// Creates a splicing iterator that replaces the specified range in the deque with the given + /// `replace_with` iterator and yields the removed items. `replace_with` does not need to be the + /// same length as `range`. + /// + /// `range` is removed even if the `Splice` iterator is not consumed before it is dropped. + /// + /// It is unspecified how many elements are removed from the deque if the `Splice` value is + /// leaked. + /// + /// The input iterator `replace_with` is only consumed when the `Splice` value is dropped. + /// + /// This is optimal if: + /// + /// * The tail (elements in the deque after `range`) is empty, + /// * or `replace_with` yields fewer or equal elements than `range`'s length + /// * or the lower bound of its `size_hint()` is exact. + /// + /// Otherwise, a temporary vector is allocated and the tail is moved twice. + /// + /// # Panics + /// + /// Panics if the range has `start_bound > end_bound`, or, if the range is + /// bounded on either end and past the length of the deque. + /// + /// # Examples + /// + /// ``` + /// # #![feature(deque_extend_front)] + /// # use std::collections::VecDeque; + /// + /// let mut v = VecDeque::from(vec![1, 2, 3, 4]); + /// let new = [7, 8, 9]; + /// let u: Vec<_> = v.splice(1..3, new).collect(); + /// assert_eq!(v, [1, 7, 8, 9, 4]); + /// assert_eq!(u, [2, 3]); + /// ``` + /// + /// Using `splice` to insert new items into a vector efficiently at a specific position + /// indicated by an empty range: + /// + /// ``` + /// # #![feature(deque_extend_front)] + /// # use std::collections::VecDeque; + /// + /// let mut v = VecDeque::from(vec![1, 5]); + /// let new = [2, 3, 4]; + /// v.splice(1..1, new); + /// assert_eq!(v, [1, 2, 3, 4, 5]); + /// ``` + #[unstable(feature = "deque_extend_front", issue = "146975")] + pub fn splice(&mut self, range: R, replace_with: I) -> Splice<'_, I::IntoIter, A> + where + R: RangeBounds, + I: IntoIterator, + { + Splice { drain: self.drain(range), replace_with: replace_with.into_iter() } + } + /// Clears the deque, removing all values. /// /// # Examples diff --git a/library/alloc/src/collections/vec_deque/splice.rs b/library/alloc/src/collections/vec_deque/splice.rs new file mode 100644 index 0000000000000..e9b2ab1c47b04 --- /dev/null +++ b/library/alloc/src/collections/vec_deque/splice.rs @@ -0,0 +1,143 @@ +use core::alloc::Allocator; + +use crate::alloc::Global; +use crate::collections::vec_deque::Drain; +use crate::vec::Vec; + +/// A splicing iterator for `VecDeque`. +/// +/// This struct is created by [`VecDeque::splice()`][super::VecDeque::splice]. +/// See its documentation for more. +/// +/// # Example +/// +/// ``` +/// # #![feature(deque_extend_front)] +/// # use std::collections::VecDeque; +/// +/// let mut v = VecDeque::from(vec![0, 1, 2]); +/// let new = [7, 8]; +/// let iter: std::collections::vec_deque::Splice<'_, _> = v.splice(1.., new); +/// ``` +#[unstable(feature = "deque_extend_front", issue = "146975")] +#[derive(Debug)] +pub struct Splice< + 'a, + I: Iterator + 'a, + #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + 'a = Global, +> { + pub(super) drain: Drain<'a, I::Item, A>, + pub(super) replace_with: I, +} + +#[unstable(feature = "deque_extend_front", issue = "146975")] +impl Iterator for Splice<'_, I, A> { + type Item = I::Item; + + fn next(&mut self) -> Option { + self.drain.next() + } + + fn size_hint(&self) -> (usize, Option) { + self.drain.size_hint() + } +} + +#[unstable(feature = "deque_extend_front", issue = "146975")] +impl DoubleEndedIterator for Splice<'_, I, A> { + fn next_back(&mut self) -> Option { + self.drain.next_back() + } +} + +#[unstable(feature = "deque_extend_front", issue = "146975")] +impl ExactSizeIterator for Splice<'_, I, A> {} + +#[unstable(feature = "deque_extend_front", issue = "146975")] +impl Drop for Splice<'_, I, A> { + fn drop(&mut self) { + // This will set drain.remaining to 0, so its drop won't try to read deallocated memory on + // drop. + self.drain.by_ref().for_each(drop); + + // At this point draining is done and the only remaining tasks are splicing + // and moving things into the final place. + + unsafe { + let source_deque = self.drain.deque.as_mut(); + + let tail_len = self.drain.tail_len; // #elements behind the drain + + if tail_len == 0 { + source_deque.extend(self.replace_with.by_ref()); + return; + } + + if !self.drain.fill(&mut self.replace_with) { + return; + } + + // There may be more elements. Use the lower bound as an estimate. + // FIXME: Is the upper bound a better guess? Or something else? + let (lower_bound, _upper_bound) = self.replace_with.size_hint(); + if lower_bound > 0 { + self.drain.move_tail(lower_bound); + if true { + // if !self.drain.fill(tail_start, &mut self.replace_with) { + self.drain.fill(&mut self.replace_with); + return; + } + } + + // Collect any remaining elements. + // This is a zero-length vector which does not allocate if `lower_bound` was exact. + let mut collected = self.replace_with.by_ref().collect::>().into_iter(); + // Now we have an exact count. + if collected.len() > 0 { + self.drain.move_tail(collected.len()); + let filled = self.drain.fill(&mut collected); + debug_assert!(filled); + debug_assert_eq!(collected.len(), 0); + } + } + // Let `Drain::drop` move the tail back if necessary and restore `deque.len`. + } +} + +/// Private helper methods for `Splice::drop` +impl Drain<'_, T, A> { + /// The range from `self.deque.len` to `self.deque.len + self.drain_len` contains elements that + /// have been moved out. + /// Fill that range as much as possible with new elements from the `replace_with` iterator. + /// Returns `true` if we filled the entire range. (`replace_with.next()` didn’t return `None`.) + unsafe fn fill>(&mut self, replace_with: &mut I) -> bool { + let deque = unsafe { self.deque.as_mut() }; + let range_start = deque.len; + let range_end = range_start + self.drain_len; + + for idx in range_start..range_end { + if let Some(new_item) = replace_with.next() { + let index = unsafe { self.deque.as_ref() }.to_physical_idx(idx); + unsafe { self.deque.as_mut().buffer_write(index, new_item) }; + deque.len += 1; + self.drain_len -= 1; + } else { + return false; + } + } + true + } + + /// Makes room for inserting more elements before the tail. + unsafe fn move_tail(&mut self, additional: usize) { + let deque = unsafe { self.deque.as_mut() }; + let tail_start = deque.len + self.drain_len; + deque.buf.reserve(tail_start + self.tail_len, additional); + + let new_tail_start = tail_start + additional; + unsafe { + deque.wrap_copy(tail_start, new_tail_start, self.tail_len); + } + self.drain_len += additional; + } +} diff --git a/library/alloctests/tests/lib.rs b/library/alloctests/tests/lib.rs index 49fb21ef5f3ac..f4f11a0a85dce 100644 --- a/library/alloctests/tests/lib.rs +++ b/library/alloctests/tests/lib.rs @@ -1,5 +1,6 @@ #![feature(allocator_api)] #![feature(alloc_layout_extra)] +#![feature(deque_extend_front)] #![feature(iter_array_chunks)] #![feature(assert_matches)] #![feature(wtf8_internals)] diff --git a/library/alloctests/tests/vec_deque.rs b/library/alloctests/tests/vec_deque.rs index a82906d55e5d0..16b531fa1efb7 100644 --- a/library/alloctests/tests/vec_deque.rs +++ b/library/alloctests/tests/vec_deque.rs @@ -1849,3 +1849,74 @@ fn test_truncate_front() { v.truncate_front(5); assert_eq!(v.as_slices(), ([2, 3, 4, 5, 6].as_slice(), [].as_slice())); } + +#[test] +fn test_splice() { + let mut v = VecDeque::from(vec![1, 2, 3, 4, 5]); + let a = [10, 11, 12]; + v.splice(2..4, a); + assert_eq!(v, &[1, 2, 10, 11, 12, 5]); + v.splice(1..3, Some(20)); + assert_eq!(v, &[1, 20, 11, 12, 5]); +} + +#[test] +fn test_splice_inclusive_range() { + let mut v = VecDeque::from(vec![1, 2, 3, 4, 5]); + let a = [10, 11, 12]; + let t1: Vec<_> = v.splice(2..=3, a).collect(); + assert_eq!(v, &[1, 2, 10, 11, 12, 5]); + assert_eq!(t1, &[3, 4]); + let t2: Vec<_> = v.splice(1..=2, Some(20)).collect(); + assert_eq!(v, &[1, 20, 11, 12, 5]); + assert_eq!(t2, &[2, 10]); +} + +#[test] +fn test_splice_inclusive_range2() { + let mut v = VecDeque::from(vec![1, 2, 10, 11, 12, 5]); + let t2: Vec<_> = v.splice(1..=2, Some(20)).collect(); + assert_eq!(v, &[1, 20, 11, 12, 5]); + assert_eq!(t2, &[2, 10]); +} + +#[test] +#[should_panic] +fn test_splice_out_of_bounds() { + let mut v = VecDeque::from(vec![1, 2, 3, 4, 5]); + let a = [10, 11, 12]; + v.splice(5..6, a); +} + +#[test] +#[should_panic] +fn test_splice_inclusive_out_of_bounds() { + let mut v = VecDeque::from(vec![1, 2, 3, 4, 5]); + let a = [10, 11, 12]; + v.splice(5..=5, a); +} + +#[test] +fn test_splice_items_zero_sized() { + let mut vec = VecDeque::from(vec![(), (), ()]); + let vec2 = VecDeque::from(vec![]); + let t: Vec<_> = vec.splice(1..2, vec2.iter().cloned()).collect(); + assert_eq!(vec, &[(), ()]); + assert_eq!(t, &[()]); +} + +#[test] +fn test_splice_unbounded() { + let mut vec = VecDeque::from(vec![1, 2, 3, 4, 5]); + let t: Vec<_> = vec.splice(.., None).collect(); + assert_eq!(vec, &[]); + assert_eq!(t, &[1, 2, 3, 4, 5]); +} + +#[test] +fn test_splice_forget() { + let mut v = VecDeque::from(vec![1, 2, 3, 4, 5]); + let a = [10, 11, 12]; + std::mem::forget(v.splice(2..4, a)); + assert_eq!(v, &[1, 2]); +}