Auto merge of #150595 - antonilol:deque_extend_front_prepend_drain_spec, r=jhpratt

Add specialization for `deque1.prepend(deque2.drain(range))` (VecDeque::prepend and extend_front)

Tracking issue: #146975

This specialization makes sure `deque1.prepend(deque2.drain(..))` gets similar speed to `deque1.append(&mut deque2)`. `deque1.prepend(deque2.drain(..))` is the equivalent of `deque1.append(&mut deque2)` but appending to the front (see the [second example of prepend](https://doc.rust-lang.org/nightly/std/collections/struct.VecDeque.html#method.prepend), prepend is from the same tracking issue).

Author: bors

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();
 

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<T, A: Allocator> SpecExtendFront<T, vec::IntoIter<T>> for VecDeque<T, A> {
     #[track_caller]
     fn spec_extend_front(&mut self, mut iterator: vec::IntoIter<T>) {
         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<T, A: Allocator> SpecExtendFront<T, Rev<vec::IntoIter<T>>> for VecDeque<T,
     fn spec_extend_front(&mut self, iterator: Rev<vec::IntoIter<T>>) {
         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<slice::Iter<'a, T>>) {
         let slice = iter.into_inner().as_slice();
-        // SAFETY: T is Copy because Copied<slice::Iter<'a, T>> is Iterator
+        self.reserve(slice.len());
+        // SAFETY: `slice.len()` space was just reserved and T is Copy because Copied<slice::Iter<'a, T>> is Iterator
         unsafe { prepend_reversed(self, slice) };
     }
 }
@@ -194,30 +197,83 @@ where
     #[track_caller]
     fn spec_extend_front(&mut self, iter: Rev<Copied<slice::Iter<'a, T>>>) {
         let slice = iter.into_inner().into_inner().as_slice();
-        // SAFETY: T is Copy because Rev<Copied<slice::Iter<'a, T>>> is Iterator
+        self.reserve(slice.len());
+        // SAFETY: `slice.len()` space was just reserved and T is Copy because Rev<Copied<slice::Iter<'a, T>>> is Iterator
         unsafe { prepend(self, slice) };
     }
 }
 
+impl<'a, T, A1: Allocator, A2: Allocator> SpecExtendFront<T, Drain<'a, T, A2>> for VecDeque<T, A1> {
+    #[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<T, Rev<Drain<'a, T, A2>>>
+    for VecDeque<T, A1>
+{
+    #[track_caller]
+    fn spec_extend_front(&mut self, iter: Rev<Drain<'a, T, A2>>) {
+        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<T, A: Allocator>(deque: &mut VecDeque<T, A>, slice: &[T]) {
-    deque.reserve(slice.len());
-
     unsafe {
         deque.head = deque.wrap_sub(deque.head, slice.len());
         deque.copy_slice(deque.head, slice);
         deque.len += slice.len();
     }
 }
 
+/// 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<T, A: Allocator>(deque: &mut VecDeque<T, A>, slice: &[T]) {
-    deque.reserve(slice.len());
-
     unsafe {
         deque.head = deque.wrap_sub(deque.head, slice.len());
         deque.copy_slice_reversed(deque.head, slice);

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<i32> {
+        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]);