Add VecDeque::extend from TrustedLen specialization

3 files changed, 175 insertions, 0 deletions

diff --git a/library/alloc/src/collections/vec_deque/mod.rs b/library/alloc/src/collections/vec_deque/mod.rs
index e28a94386c7..f92e5759b6f 100644
--- a/library/alloc/src/collections/vec_deque/mod.rs
+++ b/library/alloc/src/collections/vec_deque/mod.rs
@@ -453,6 +453,25 @@ impl<T, A: Allocator> VecDeque<T, A> {
         }
     }
 
+    /// Writes all values from `iter` to `dst`.
+    ///
+    /// # Safety
+    ///
+    /// Assumes no wrapping around happens.
+    /// Assumes capacity is sufficient.
+    #[inline]
+    unsafe fn write_iter(
+        &mut self,
+        dst: usize,
+        iter: impl Iterator<Item = T>,
+        written: &mut usize,
+    ) {
+        iter.enumerate().for_each(|(i, element)| unsafe {
+            self.buffer_write(dst + i, element);
+            *written += 1;
+        });
+    }
+
     /// Frobs the head and tail sections around to handle the fact that we
     /// just reallocated. Unsafe because it trusts old_capacity.
     #[inline]
diff --git a/library/alloc/src/collections/vec_deque/spec_extend.rs b/library/alloc/src/collections/vec_deque/spec_extend.rs
index 172f2e9068f..cea2b7543ce 100644
--- a/library/alloc/src/collections/vec_deque/spec_extend.rs
+++ b/library/alloc/src/collections/vec_deque/spec_extend.rs
@@ -1,5 +1,6 @@
 use crate::alloc::Allocator;
 use crate::vec;
+use core::iter::TrustedLen;
 use core::slice;
 
 use super::VecDeque;
@@ -34,6 +35,64 @@ where
     }
 }
 
+impl<T, I, A: Allocator> SpecExtend<T, I> for VecDeque<T, A>
+where
+    I: TrustedLen<Item = T>,
+{
+    default fn spec_extend(&mut self, mut iter: I) {
+        // This is the case for a TrustedLen iterator.
+        let (low, high) = iter.size_hint();
+        if let Some(additional) = high {
+            debug_assert_eq!(
+                low,
+                additional,
+                "TrustedLen iterator's size hint is not exact: {:?}",
+                (low, high)
+            );
+            self.reserve(additional);
+
+            struct WrapAddOnDrop<'a, T, A: Allocator> {
+                vec_deque: &'a mut VecDeque<T, A>,
+                written: usize,
+            }
+
+            impl<'a, T, A: Allocator> Drop for WrapAddOnDrop<'a, T, A> {
+                fn drop(&mut self) {
+                    self.vec_deque.head =
+                        self.vec_deque.wrap_add(self.vec_deque.head, self.written);
+                }
+            }
+
+            let mut wrapper = WrapAddOnDrop { vec_deque: self, written: 0 };
+
+            let head_room = wrapper.vec_deque.cap() - wrapper.vec_deque.head;
+            unsafe {
+                wrapper.vec_deque.write_iter(
+                    wrapper.vec_deque.head,
+                    ByRefSized(&mut iter).take(head_room),
+                    &mut wrapper.written,
+                );
+
+                if additional > head_room {
+                    wrapper.vec_deque.write_iter(0, iter, &mut wrapper.written);
+                }
+            }
+
+            debug_assert_eq!(
+                additional, wrapper.written,
+                "The number of items written to VecDeque doesn't match the TrustedLen size hint"
+            );
+        } else {
+            // Per TrustedLen contract a `None` upper bound means that the iterator length
+            // truly exceeds usize::MAX, which would eventually lead to a capacity overflow anyway.
+            // Since the other branch already panics eagerly (via `reserve()`) we do the same here.
+            // This avoids additional codegen for a fallback code path which would eventually
+            // panic anyway.
+            panic!("capacity overflow");
+        }
+    }
+}
+
 impl<T, A: Allocator> SpecExtend<T, vec::IntoIter<T>> for VecDeque<T, A> {
     fn spec_extend(&mut self, mut iterator: vec::IntoIter<T>) {
         let slice = iterator.as_slice();
diff --git a/library/alloc/src/collections/vec_deque/tests.rs b/library/alloc/src/collections/vec_deque/tests.rs
index f2869a54713..1f2daef213c 100644
--- a/library/alloc/src/collections/vec_deque/tests.rs
+++ b/library/alloc/src/collections/vec_deque/tests.rs
@@ -1,3 +1,5 @@
+use core::iter::TrustedLen;
+
 use super::*;
 
 #[bench]
@@ -792,6 +794,101 @@ fn test_from_vec() {
 }
 
 #[test]
+fn test_extend_basic() {
+    test_extend_impl(false);
+}
+
+#[test]
+fn test_extend_trusted_len() {
+    test_extend_impl(true);
+}
+
+fn test_extend_impl(trusted_len: bool) {
+    struct VecDequeTester {
+        test: VecDeque<usize>,
+        expected: VecDeque<usize>,
+        trusted_len: bool,
+    }
+
+    impl VecDequeTester {
+        fn new(trusted_len: bool) -> Self {
+            Self { test: VecDeque::new(), expected: VecDeque::new(), trusted_len }
+        }
+
+        fn test_extend<I>(&mut self, iter: I)
+        where
+            I: Iterator<Item = usize> + TrustedLen + Clone,
+        {
+            struct BasicIterator<I>(I);
+            impl<I> Iterator for BasicIterator<I>
+            where
+                I: Iterator<Item = usize>,
+            {
+                type Item = usize;
+
+                fn next(&mut self) -> Option<Self::Item> {
+                    self.0.next()
+                }
+            }
+
+            if self.trusted_len {
+                self.test.extend(iter.clone());
+            } else {
+                self.test.extend(BasicIterator(iter.clone()));
+            }
+
+            for item in iter {
+                self.expected.push_back(item)
+            }
+
+            assert_eq!(self.test, self.expected);
+            let (a1, b1) = self.test.as_slices();
+            let (a2, b2) = self.expected.as_slices();
+            assert_eq!(a1, a2);
+            assert_eq!(b1, b2);
+        }
+
+        fn drain<R: RangeBounds<usize> + Clone>(&mut self, range: R) {
+            self.test.drain(range.clone());
+            self.expected.drain(range);
+
+            assert_eq!(self.test, self.expected);
+        }
+
+        fn clear(&mut self) {
+            self.test.clear();
+            self.expected.clear();
+        }
+
+        fn remaining_capacity(&self) -> usize {
+            self.test.capacity() - self.test.len()
+        }
+    }
+
+    let mut tester = VecDequeTester::new(trusted_len);
+
+    // Initial capacity
+    tester.test_extend(0..tester.remaining_capacity() - 1);
+
+    // Grow
+    tester.test_extend(1024..2048);
+
+    // Wrap around
+    tester.drain(..128);
+
+    tester.test_extend(0..tester.remaining_capacity() - 1);
+
+    // Continue
+    tester.drain(256..);
+    tester.test_extend(4096..8196);
+
+    tester.clear();
+
+    // Start again
+    tester.test_extend(0..32);
+}
+
+#[test]
 #[should_panic = "capacity overflow"]
 fn test_from_vec_zst_overflow() {
     use crate::vec::Vec;