Add VecDeque::extend from TrustedLen specialization

library/alloc/benches/vec_deque.rs

@@ -91,3 +91,35 @@ fn bench_extend_vec(b: &mut Bencher) {
  ring.extend(black_box(input));
  });
 }
+
+#[bench]
+fn bench_extend_trustedlen(b: &mut Bencher) {
+ let mut ring: VecDeque<u16> = VecDeque::with_capacity(1000);
+
+ b.iter(|| {
+ ring.clear();
+ ring.extend(black_box(0..512));
+ });
+}
+
+#[bench]
+fn bench_extend_chained_trustedlen(b: &mut Bencher) {
+ let mut ring: VecDeque<u16> = VecDeque::with_capacity(1000);
+
+ b.iter(|| {
+ ring.clear();
+ ring.extend(black_box((0..256).chain(768..1024)));
+ });
+}
+
+#[bench]
+fn bench_extend_chained_bytes(b: &mut Bencher) {
+ let mut ring: VecDeque<u16> = VecDeque::with_capacity(1000);
+ let input1: &[u16] = &[128; 256];
+ let input2: &[u16] = &[255; 256];
+
+ b.iter(|| {
+ ring.clear();
+ ring.extend(black_box(input1.iter().chain(input2.iter())));
+ });
+}

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]

library/alloc/src/collections/vec_deque/spec_extend.rs

@@ -1,5 +1,6 @@
 use crate::alloc::Allocator;
 use crate::vec;
+use core::iter::{ByRefSized, 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();

library/alloc/src/collections/vec_deque/tests.rs

@@ -1,3 +1,5 @@
+use core::iter::TrustedLen;
+
 use super::*;
 
 #[bench]
@@ -791,6 +793,101 @@ fn test_from_vec() {
  assert_eq!(vd.len(), vec.len());
 }
 
+#[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() {

library/alloc/src/lib.rs

@@ -143,6 +143,7 @@
 #![feature(unchecked_math)]
 #![feature(unicode_internals)]
 #![feature(unsize)]
+#![feature(std_internals)]
 //
 // Language features:
 #![feature(allocator_internals)]

library/core/src/iter/adapters/by_ref_sized.rs

@@ -4,8 +4,11 @@ use crate::ops::Try;
 ///
 /// Ideally this will no longer be required, eventually, but as can be seen in
 /// the benchmarks (as of Feb 2022 at least) `by_ref` can have performance cost.
-pub(crate) struct ByRefSized<'a, I>(pub &'a mut I);
+#[unstable(feature = "std_internals", issue = "none")]
+#[derive(Debug)]
+pub struct ByRefSized<'a, I>(pub &'a mut I);
 
+#[unstable(feature = "std_internals", issue = "none")]
 impl<I: Iterator> Iterator for ByRefSized<'_, I> {
  type Item = I::Item;
 
@@ -47,6 +50,7 @@ impl<I: Iterator> Iterator for ByRefSized<'_, I> {
  }
 }
 
+#[unstable(feature = "std_internals", issue = "none")]
 impl<I: DoubleEndedIterator> DoubleEndedIterator for ByRefSized<'_, I> {
  #[inline]
  fn next_back(&mut self) -> Option<Self::Item> {

library/core/src/iter/adapters/mod.rs

@@ -32,7 +32,8 @@ pub use self::{
  scan::Scan, skip::Skip, skip_while::SkipWhile, take::Take, take_while::TakeWhile, zip::Zip,
 };
 
-pub(crate) use self::by_ref_sized::ByRefSized;
+#[unstable(feature = "std_internals", issue = "none")]
+pub use self::by_ref_sized::ByRefSized;
 
 #[stable(feature = "iter_cloned", since = "1.1.0")]
 pub use self::cloned::Cloned;

library/core/src/iter/mod.rs

@@ -398,6 +398,8 @@ pub use self::traits::{
 
 #[stable(feature = "iter_zip", since = "1.59.0")]
 pub use self::adapters::zip;
+#[unstable(feature = "std_internals", issue = "none")]
+pub use self::adapters::ByRefSized;
 #[stable(feature = "iter_cloned", since = "1.1.0")]
 pub use self::adapters::Cloned;
 #[stable(feature = "iter_copied", since = "1.36.0")]
@@ -422,7 +424,7 @@ pub use self::adapters::{
 #[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
 pub use self::adapters::{Intersperse, IntersperseWith};
 
-pub(crate) use self::adapters::{try_process, ByRefSized};
+pub(crate) use self::adapters::try_process;
 
 mod adapters;
 mod range;