Make VecDeque::from_iter O(1) from vec(_deque)::IntoIter

Author: scottmcm

alloc/src/collections/vec_deque/into_iter.rs

@@ -25,6 +25,10 @@ impl<T, A: Allocator> IntoIter<T, A> {
     pub(super) fn new(inner: VecDeque<T, A>) -> Self {
         IntoIter { inner }
     }
+
+    pub(super) fn into_vecdeque(self) -> VecDeque<T, A> {
+        self.inner
+    }
 }
 
 #[stable(feature = "collection_debug", since = "1.17.0")]

alloc/src/collections/vec_deque/mod.rs

@@ -55,6 +55,10 @@ use self::spec_extend::SpecExtend;
 
 mod spec_extend;
 
+use self::spec_from_iter::SpecFromIter;
+
+mod spec_from_iter;
+
 #[cfg(test)]
 mod tests;
 
@@ -586,6 +590,35 @@ impl<T, A: Allocator> VecDeque<T, A> {
         VecDeque { head: 0, len: 0, buf: RawVec::with_capacity_in(capacity, alloc) }
     }
 
+    /// For use by `vec::IntoIter::into_vecdeque`
+    ///
+    /// # Safety
+    ///
+    /// All the usual requirements on the allocated memory like in
+    /// `Vec::from_raw_parts_in`, but takes a *range* of elements that are
+    /// initialized rather than only supporting `0..len`.  Requires that
+    /// `initialized.start` ≤ `initialized.end` ≤ `capacity`.
+    #[inline]
+    pub(crate) unsafe fn from_contiguous_raw_parts_in(
+        ptr: *mut T,
+        initialized: Range<usize>,
+        capacity: usize,
+        alloc: A,
+    ) -> Self {
+        debug_assert!(initialized.start <= initialized.end);
+        debug_assert!(initialized.end <= capacity);
+
+        // SAFETY: Our safety precondition guarantees the range length won't wrap,
+        // and that the allocation is valid for use in `RawVec`.
+        unsafe {
+            VecDeque {
+                head: initialized.start,
+                len: initialized.end.unchecked_sub(initialized.start),
+                buf: RawVec::from_raw_parts_in(ptr, capacity, alloc),
+            }
+        }
+    }
+
     /// Provides a reference to the element at the given index.
     ///
     /// Element at index 0 is the front of the queue.
@@ -2699,18 +2732,8 @@ impl<T, A: Allocator> IndexMut<usize> for VecDeque<T, A> {
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<T> FromIterator<T> for VecDeque<T> {
-    #[inline]
     fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> VecDeque<T> {
-        // Since converting is O(1) now, might as well re-use that logic
-        // (including things like the `vec::IntoIter`→`Vec` specialization)
-        // especially as that could save us some monomorphiziation work
-        // if one uses the same iterators (like slice ones) with both.
-        return from_iter_via_vec(iter.into_iter());
-
-        #[inline]
-        fn from_iter_via_vec<U>(iter: impl Iterator<Item = U>) -> VecDeque<U> {
-            Vec::from_iter(iter).into()
-        }
+        SpecFromIter::spec_from_iter(iter.into_iter())
     }
 }
 

alloc/src/collections/vec_deque/spec_from_iter.rs

@@ -0,0 +1,33 @@
+use super::{IntoIter, VecDeque};
+
+/// Specialization trait used for `VecDeque::from_iter`
+pub(super) trait SpecFromIter<T, I> {
+    fn spec_from_iter(iter: I) -> Self;
+}
+
+impl<T, I> SpecFromIter<T, I> for VecDeque<T>
+where
+    I: Iterator<Item = T>,
+{
+    default fn spec_from_iter(iterator: I) -> Self {
+        // Since converting is O(1) now, just re-use the `Vec` logic for
+        // anything where we can't do something extra-special for `VecDeque`,
+        // especially as that could save us some monomorphiziation work
+        // if one uses the same iterators (like slice ones) with both.
+        crate::vec::Vec::from_iter(iterator).into()
+    }
+}
+
+impl<T> SpecFromIter<T, crate::vec::IntoIter<T>> for VecDeque<T> {
+    #[inline]
+    fn spec_from_iter(iterator: crate::vec::IntoIter<T>) -> Self {
+        iterator.into_vecdeque()
+    }
+}
+
+impl<T> SpecFromIter<T, IntoIter<T>> for VecDeque<T> {
+    #[inline]
+    fn spec_from_iter(iterator: IntoIter<T>) -> Self {
+        iterator.into_vecdeque()
+    }
+}

alloc/src/vec/into_iter.rs

@@ -1,6 +1,7 @@
 #[cfg(not(no_global_oom_handling))]
 use super::AsVecIntoIter;
 use crate::alloc::{Allocator, Global};
+use crate::collections::VecDeque;
 use crate::raw_vec::RawVec;
 use core::array;
 use core::fmt;
@@ -132,6 +133,32 @@ impl<T, A: Allocator> IntoIter<T, A> {
     pub(crate) fn forget_remaining_elements(&mut self) {
         self.ptr = self.end;
     }
+
+    #[inline]
+    pub(crate) fn into_vecdeque(self) -> VecDeque<T, A> {
+        // Keep our `Drop` impl from dropping the elements and the allocator
+        let mut this = ManuallyDrop::new(self);
+
+        // SAFETY: This allocation originally came from a `Vec`, so it passes
+        // all those checks.  We have `this.buf` ≤ `this.ptr` ≤ `this.end`,
+        // so the `sub_ptr`s below cannot wrap, and will produce a well-formed
+        // range.  `end` ≤ `buf + cap`, so the range will be in-bounds.
+        // Taking `alloc` is ok because nothing else is going to look at it,
+        // since our `Drop` impl isn't going to run so there's no more code.
+        unsafe {
+            let buf = this.buf.as_ptr();
+            let initialized = if T::IS_ZST {
+                // All the pointers are the same for ZSTs, so it's fine to
+                // say that they're all at the beginning of the "allocation".
+                0..this.len()
+            } else {
+                this.ptr.sub_ptr(buf)..this.end.sub_ptr(buf)
+            };
+            let cap = this.cap;
+            let alloc = ManuallyDrop::take(&mut this.alloc);
+            VecDeque::from_contiguous_raw_parts_in(buf, initialized, cap, alloc)
+        }
+    }
 }
 
 #[stable(feature = "vec_intoiter_as_ref", since = "1.46.0")]

alloc/tests/vec_deque.rs

@@ -1736,3 +1736,29 @@ fn test_resize_keeps_reserved_space_from_item() {
     d.resize(1, v);
     assert_eq!(d[0].capacity(), 1234);
 }
+
+#[test]
+fn test_collect_from_into_iter_keeps_allocation() {
+    let mut v = Vec::with_capacity(13);
+    v.extend(0..7);
+    check(v.into_iter());
+
+    let mut v = VecDeque::with_capacity(13);
+    v.extend(0..7);
+    check(v.into_iter());
+
+    fn check(mut it: impl Iterator<Item = i32>) {
+        assert_eq!(it.next(), Some(0));
+        assert_eq!(it.next(), Some(1));
+        let mut v: VecDeque<i32> = it.collect();
+        assert_eq!(v.capacity(), 13);
+        assert_eq!(v.as_slices(), ([2, 3, 4, 5, 6].as_slice(), [].as_slice()));
+        v.push_front(7);
+        assert_eq!(v.as_slices(), ([7, 2, 3, 4, 5, 6].as_slice(), [].as_slice()));
+        v.push_front(8);
+        assert_eq!(v.as_slices(), ([8, 7, 2, 3, 4, 5, 6].as_slice(), [].as_slice()));
+        v.push_front(9);
+        assert_eq!(v.as_slices(), ([9].as_slice(), [8, 7, 2, 3, 4, 5, 6].as_slice()));
+        assert_eq!(v.capacity(), 13);
+    }
+}