From 0e2b98236b415d38b3ecf44b22a0b3eb4ce31304 Mon Sep 17 00:00:00 2001 From: lcnr/Bastian Kauschke Date: Mon, 11 Mar 2019 22:27:55 +0100 Subject: [PATCH] implement drain --- lib.rs | 134 +++++++++++++++++++++++++++++++++++++++++++++------------ 1 file changed, 107 insertions(+), 27 deletions(-) diff --git a/lib.rs b/lib.rs index 96d87bf..c9cddf4 100644 --- a/lib.rs +++ b/lib.rs @@ -60,13 +60,13 @@ use std::hash::{Hash, Hasher}; use std::hint::unreachable_unchecked; #[cfg(feature = "std")] use std::io; -use std::iter::{repeat, FromIterator, IntoIterator}; +use std::iter::{repeat, FromIterator, FusedIterator, IntoIterator}; #[cfg(feature = "serde")] use std::marker::PhantomData; use std::mem; use std::mem::MaybeUninit; -use std::ops; -use std::ptr; +use std::ops::{self, Bound, RangeBounds}; +use std::ptr::{self, NonNull}; use std::slice::{self, SliceIndex}; /// Creates a [`SmallVec`] containing the arguments. @@ -182,15 +182,30 @@ unsafe fn deallocate(ptr: *mut T, capacity: usize) { /// Returned from [`SmallVec::drain`][1]. /// /// [1]: struct.SmallVec.html#method.drain -pub struct Drain<'a, T: 'a> { - iter: slice::IterMut<'a, T>, +pub struct Drain<'a, T: 'a + Array> { + tail_start: usize, + tail_len: usize, + iter: slice::Iter<'a, T::Item>, + vec: NonNull>, } -impl<'a, T: 'a> Iterator for Drain<'a, T> { - type Item = T; +impl<'a, T: 'a + Array> fmt::Debug for Drain<'a, T> +where + T::Item: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_tuple("Drain").field(&self.iter.as_slice()).finish() + } +} + +unsafe impl<'a, T: Sync + Array> Sync for Drain<'a, T> {} +unsafe impl<'a, T: Send + Array> Send for Drain<'a, T> {} + +impl<'a, T: 'a + Array> Iterator for Drain<'a, T> { + type Item = T::Item; #[inline] - fn next(&mut self) -> Option { + fn next(&mut self) -> Option { self.iter .next() .map(|reference| unsafe { ptr::read(reference) }) @@ -202,21 +217,43 @@ impl<'a, T: 'a> Iterator for Drain<'a, T> { } } -impl<'a, T: 'a> DoubleEndedIterator for Drain<'a, T> { +impl<'a, T: 'a + Array> DoubleEndedIterator for Drain<'a, T> { #[inline] - fn next_back(&mut self) -> Option { + fn next_back(&mut self) -> Option { self.iter .next_back() .map(|reference| unsafe { ptr::read(reference) }) } } -impl<'a, T> ExactSizeIterator for Drain<'a, T> {} +impl<'a, T: Array> ExactSizeIterator for Drain<'a, T> { + #[inline] + fn len(&self) -> usize { + self.iter.len() + } +} -impl<'a, T: 'a> Drop for Drain<'a, T> { +impl<'a, T: Array> FusedIterator for Drain<'a, T> {} + +impl<'a, T: 'a + Array> Drop for Drain<'a, T> { fn drop(&mut self) { - // Destroy the remaining elements. - for _ in self.by_ref() {} + self.for_each(drop); + + if self.tail_len > 0 { + unsafe { + let source_vec = self.vec.as_mut(); + + // memmove back untouched tail, update to new length + let start = source_vec.len(); + let tail = self.tail_start; + if tail != start { + let src = source_vec.as_ptr().add(tail); + let dst = source_vec.as_mut_ptr().add(start); + ptr::copy(src, dst, self.tail_len); + } + source_vec.set_len(start + self.tail_len); + } + } } } @@ -548,18 +585,50 @@ impl SmallVec { self.capacity > A::size() } - /// Empty the vector and return an iterator over its former contents. - pub fn drain(&mut self) -> Drain { - unsafe { - let ptr = self.as_mut_ptr(); + /// Creates a draining iterator that removes the specified range in the vector + /// and yields the removed items. + /// + /// Note 1: The element range is removed even if the iterator is only + /// partially consumed or not consumed at all. + /// + /// Note 2: It is unspecified how many elements are removed from the vector + /// if the `Drain` value is leaked. + /// + /// # Panics + /// + /// Panics if the starting point is greater than the end point or if + /// the end point is greater than the length of the vector. + pub fn drain(&mut self, range: R) -> Drain + where + R: RangeBounds, + { + use Bound::*; - let current_len = self.len(); - self.set_len(0); + let len = self.len(); + let start = match range.start_bound() { + Included(&n) => n, + Excluded(&n) => n + 1, + Unbounded => 0, + }; + let end = match range.end_bound() { + Included(&n) => n + 1, + Excluded(&n) => n, + Unbounded => len, + }; + + assert!(start <= end); + assert!(end <= len); + + unsafe { + self.set_len(start); - let slice = slice::from_raw_parts_mut(ptr, current_len); + let range_slice = slice::from_raw_parts_mut(self.as_mut_ptr().add(start), end - start); Drain { - iter: slice.iter_mut(), + tail_start: end, + tail_len: len - end, + iter: range_slice.iter(), + vec: NonNull::from(self), } } } @@ -1697,26 +1766,36 @@ mod tests { fn drain() { let mut v: SmallVec<[u8; 2]> = SmallVec::new(); v.push(3); - assert_eq!(v.drain().collect::>(), &[3]); + assert_eq!(v.drain(..).collect::>(), &[3]); // spilling the vec v.push(3); v.push(4); v.push(5); - assert_eq!(v.drain().collect::>(), &[3, 4, 5]); + let old_capacity = v.capacity(); + assert_eq!(v.drain(1..).collect::>(), &[4, 5]); + // drain should not change the capacity + assert_eq!(v.capacity(), old_capacity); } #[test] fn drain_rev() { let mut v: SmallVec<[u8; 2]> = SmallVec::new(); v.push(3); - assert_eq!(v.drain().rev().collect::>(), &[3]); + assert_eq!(v.drain(..).rev().collect::>(), &[3]); // spilling the vec v.push(3); v.push(4); v.push(5); - assert_eq!(v.drain().rev().collect::>(), &[5, 4, 3]); + assert_eq!(v.drain(..).rev().collect::>(), &[5, 4, 3]); + } + + #[test] + fn drain_forget() { + let mut v: SmallVec<[u8; 1]> = smallvec![0, 1, 2, 3, 4, 5, 6, 7]; + std::mem::forget(v.drain(2..5)); + assert_eq!(v.len(), 2); } #[test] @@ -2157,7 +2236,8 @@ mod tests { fn test_exact_size_iterator() { let mut vec = SmallVec::<[u32; 2]>::from(&[1, 2, 3][..]); assert_eq!(vec.clone().into_iter().len(), 3); - assert_eq!(vec.drain().len(), 3); + assert_eq!(vec.drain(..2).len(), 2); + assert_eq!(vec.into_iter().len(), 1); } #[test]