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 VecDeque { } } + /// 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, + 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 SpecExtend for VecDeque +where + I: TrustedLen, +{ + 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, + 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 SpecExtend> for VecDeque { fn spec_extend(&mut self, mut iterator: vec::IntoIter) { 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] @@ -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, + expected: VecDeque, + trusted_len: bool, + } + + impl VecDequeTester { + fn new(trusted_len: bool) -> Self { + Self { test: VecDeque::new(), expected: VecDeque::new(), trusted_len } + } + + fn test_extend(&mut self, iter: I) + where + I: Iterator + TrustedLen + Clone, + { + struct BasicIterator(I); + impl Iterator for BasicIterator + where + I: Iterator, + { + type Item = usize; + + fn next(&mut self) -> Option { + 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 + 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() {