diff options
Diffstat (limited to '')
-rw-r--r-- | library/alloc/src/collections/vec_deque/spec_extend.rs | 81 |
1 files changed, 36 insertions, 45 deletions
diff --git a/library/alloc/src/collections/vec_deque/spec_extend.rs b/library/alloc/src/collections/vec_deque/spec_extend.rs index 97ff8b765..dccf40ccb 100644 --- a/library/alloc/src/collections/vec_deque/spec_extend.rs +++ b/library/alloc/src/collections/vec_deque/spec_extend.rs @@ -1,6 +1,6 @@ use crate::alloc::Allocator; use crate::vec; -use core::iter::{ByRefSized, TrustedLen}; +use core::iter::TrustedLen; use core::slice; use super::VecDeque; @@ -17,19 +17,33 @@ where default fn spec_extend(&mut self, mut iter: I) { // This function should be the moral equivalent of: // - // for item in iter { - // self.push_back(item); - // } - while let Some(element) = iter.next() { - if self.len() == self.capacity() { - let (lower, _) = iter.size_hint(); - self.reserve(lower.saturating_add(1)); - } + // for item in iter { + // self.push_back(item); + // } + + // May only be called if `deque.len() < deque.capacity()` + unsafe fn push_unchecked<T, A: Allocator>(deque: &mut VecDeque<T, A>, element: T) { + // SAFETY: Because of the precondition, it's guaranteed that there is space + // in the logical array after the last element. + unsafe { deque.buffer_write(deque.to_physical_idx(deque.len), element) }; + // This can't overflow because `deque.len() < deque.capacity() <= usize::MAX`. + deque.len += 1; + } - let head = self.head; - self.head = self.wrap_add(self.head, 1); - unsafe { - self.buffer_write(head, element); + while let Some(element) = iter.next() { + let (lower, _) = iter.size_hint(); + self.reserve(lower.saturating_add(1)); + + // SAFETY: We just reserved space for at least one element. + unsafe { push_unchecked(self, element) }; + + // Inner loop to avoid repeatedly calling `reserve`. + while self.len < self.capacity() { + let Some(element) = iter.next() else { + return; + }; + // SAFETY: The loop condition guarantees that `self.len() < self.capacity()`. + unsafe { push_unchecked(self, element) }; } } } @@ -39,7 +53,7 @@ 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) { + default fn spec_extend(&mut self, iter: I) { // This is the case for a TrustedLen iterator. let (low, high) = iter.size_hint(); if let Some(additional) = high { @@ -51,35 +65,12 @@ where ); 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); - } - } + let written = unsafe { + self.write_iter_wrapping(self.to_physical_idx(self.len), iter, additional) + }; debug_assert_eq!( - additional, wrapper.written, + additional, written, "The number of items written to VecDeque doesn't match the TrustedLen size hint" ); } else { @@ -99,8 +90,8 @@ impl<T, A: Allocator> SpecExtend<T, vec::IntoIter<T>> for VecDeque<T, A> { self.reserve(slice.len()); unsafe { - self.copy_slice(self.head, slice); - self.head = self.wrap_add(self.head, slice.len()); + self.copy_slice(self.to_physical_idx(self.len), slice); + self.len += slice.len(); } iterator.forget_remaining_elements(); } @@ -125,8 +116,8 @@ where self.reserve(slice.len()); unsafe { - self.copy_slice(self.head, slice); - self.head = self.wrap_add(self.head, slice.len()); + self.copy_slice(self.to_physical_idx(self.len), slice); + self.len += slice.len(); } } } |