diff options
Diffstat (limited to 'src/vec.rs')
-rw-r--r-- | src/vec.rs | 68 |
1 files changed, 39 insertions, 29 deletions
@@ -138,16 +138,10 @@ impl<'data, T: Send> IndexedParallelIterator for Drain<'data, T> { { unsafe { // Make the vector forget about the drained items, and temporarily the tail too. - let start = self.range.start; - self.vec.set_len(start); + self.vec.set_len(self.range.start); // Create the producer as the exclusive "owner" of the slice. - let producer = { - // Get a correct borrow lifetime, then extend it to the original length. - let mut slice = &mut self.vec[start..]; - slice = slice::from_raw_parts_mut(slice.as_mut_ptr(), self.range.len()); - DrainProducer::new(slice) - }; + let producer = DrainProducer::from_vec(self.vec, self.range.len()); // The producer will move or drop each item from the drained range. callback.callback(producer) @@ -157,22 +151,24 @@ impl<'data, T: Send> IndexedParallelIterator for Drain<'data, T> { impl<'data, T: Send> Drop for Drain<'data, T> { fn drop(&mut self) { - if self.range.len() > 0 { - let Range { start, end } = self.range; - if self.vec.len() != start { - // We must not have produced, so just call a normal drain to remove the items. - assert_eq!(self.vec.len(), self.orig_len); - self.vec.drain(start..end); - } else if end < self.orig_len { - // The producer was responsible for consuming the drained items. - // Move the tail items to their new place, then set the length to include them. - unsafe { - let ptr = self.vec.as_mut_ptr().add(start); - let tail_ptr = self.vec.as_ptr().add(end); - let tail_len = self.orig_len - end; - ptr::copy(tail_ptr, ptr, tail_len); - self.vec.set_len(start + tail_len); - } + let Range { start, end } = self.range; + if self.vec.len() == self.orig_len { + // We must not have produced, so just call a normal drain to remove the items. + self.vec.drain(start..end); + } else if start == end { + // Empty range, so just restore the length to its original state + unsafe { + self.vec.set_len(self.orig_len); + } + } else if end < self.orig_len { + // The producer was responsible for consuming the drained items. + // Move the tail items to their new place, then set the length to include them. + unsafe { + let ptr = self.vec.as_mut_ptr().add(start); + let tail_ptr = self.vec.as_ptr().add(end); + let tail_len = self.orig_len - end; + ptr::copy(tail_ptr, ptr, tail_len); + self.vec.set_len(start + tail_len); } } } @@ -184,13 +180,27 @@ pub(crate) struct DrainProducer<'data, T: Send> { slice: &'data mut [T], } -impl<'data, T: 'data + Send> DrainProducer<'data, T> { +impl<T: Send> DrainProducer<'_, T> { /// Creates a draining producer, which *moves* items from the slice. /// - /// Unsafe bacause `!Copy` data must not be read after the borrow is released. - pub(crate) unsafe fn new(slice: &'data mut [T]) -> Self { + /// Unsafe because `!Copy` data must not be read after the borrow is released. + pub(crate) unsafe fn new(slice: &mut [T]) -> DrainProducer<'_, T> { DrainProducer { slice } } + + /// Creates a draining producer, which *moves* items from the tail of the vector. + /// + /// Unsafe because we're moving from beyond `vec.len()`, so the caller must ensure + /// that data is initialized and not read after the borrow is released. + unsafe fn from_vec(vec: &mut Vec<T>, len: usize) -> DrainProducer<'_, T> { + let start = vec.len(); + assert!(vec.capacity() - start >= len); + + // The pointer is derived from `Vec` directly, not through a `Deref`, + // so it has provenance over the whole allocation. + let ptr = vec.as_mut_ptr().add(start); + DrainProducer::new(slice::from_raw_parts_mut(ptr, len)) + } } impl<'data, T: 'data + Send> Producer for DrainProducer<'data, T> { @@ -199,7 +209,7 @@ impl<'data, T: 'data + Send> Producer for DrainProducer<'data, T> { fn into_iter(mut self) -> Self::IntoIter { // replace the slice so we don't drop it twice - let slice = mem::replace(&mut self.slice, &mut []); + let slice = mem::take(&mut self.slice); SliceDrain { iter: slice.iter_mut(), } @@ -207,7 +217,7 @@ impl<'data, T: 'data + Send> Producer for DrainProducer<'data, T> { fn split_at(mut self, index: usize) -> (Self, Self) { // replace the slice so we don't drop it twice - let slice = mem::replace(&mut self.slice, &mut []); + let slice = mem::take(&mut self.slice); let (left, right) = slice.split_at_mut(index); unsafe { (DrainProducer::new(left), DrainProducer::new(right)) } } |