diff --git a/src/drain.rs b/src/drain.rs index da00770..31fd029 100644 --- a/src/drain.rs +++ b/src/drain.rs @@ -104,17 +104,17 @@ impl Drop for Drain<'_, T> { let mut vec = self.vec; - // if T::IS_ZST { - // // ZSTs have no identity, so we don't need to move them around, we only need to drop the correct amount. - // // this can be achieved by manipulating the Vec length instead of moving values out from `iter`. - // unsafe { - // let vec = vec.as_mut(); - // let old_len = vec.len(); - // vec.set_len(old_len + drop_len + self.tail_len); - // vec.truncate(old_len + self.tail_len); - // } - // return; - // } + if mem::size_of::() == 0 { + // ZSTs have no identity, so we don't need to move them around, we only need to drop the correct amount. + // this can be achieved by manipulating the Vec length instead of moving values out from `iter`. + unsafe { + let vec = vec.as_mut(); + let old_len = vec.len(); + vec.header.len = (old_len + drop_len + self.tail_len) as u32; + vec.truncate(old_len + self.tail_len); + } + return; + }; // ensure elements are moved back into their appropriate places, even when drop_in_place panics let _guard = DropGuard(self); diff --git a/src/shared.rs b/src/shared.rs index f28f3bf..25f1d54 100644 --- a/src/shared.rs +++ b/src/shared.rs @@ -256,7 +256,7 @@ impl RefCountedVector { unsafe { let data = NonNull::new_unchecked(self.data_ptr()); let header = self.vec_header().clone(); - let allocator = self.inner.as_ref().allocator.clone(); + let allocator = ptr::read(&self.inner.as_ref().allocator); mem::forget(self); @@ -643,7 +643,7 @@ impl Drop for RefCountedVector { } -unsafe impl Send for AtomicSharedVector {} +unsafe impl Send for AtomicSharedVector {} unsafe impl Sync for AtomicSharedVector {} diff --git a/src/vector.rs b/src/vector.rs index afbaebc..7a47147 100644 --- a/src/vector.rs +++ b/src/vector.rs @@ -378,18 +378,21 @@ impl RawVector { unsafe fn extend_within_capacity(&mut self, iter: &mut impl Iterator) { let n = self.remaining_capacity() as BufferSize; + if n == 0 { + return; + } let mut ptr = self.data_ptr().add(self.len()); let mut count = 0; unsafe { for item in iter { - if count == n { - break; - } ptr::write(ptr, item); ptr = ptr.add(1); count += 1; + if count == n { + break; + } } self.header.len += count; } @@ -434,6 +437,24 @@ impl RawVector { self.try_realloc_with_capacity(allocator, new_cap) } + pub fn truncate(&mut self, new_len: usize) { + let old_len = self.header.len as usize; + if old_len <= new_len { + return; + } + + unsafe { + let mut elt = self.data_ptr().add(new_len); + let end = self.data_ptr().add(old_len); + while elt < end { + ptr::drop_in_place(elt); + elt = elt.add(1) + } + } + + self.header.len = new_len as u32; + } + #[cold] unsafe fn try_realloc_with_capacity(&mut self, allocator: &A, new_cap: usize) -> Result<(), AllocError> { type R = DefaultRefCount; @@ -447,12 +468,7 @@ impl RawVector { let old_len = self.len(); if old_len > new_cap { - let mut elt = self.data_ptr().add(new_cap); - let end = self.data_ptr().add(old_len); - while elt < end { - ptr::drop_in_place(elt); - elt = elt.add(1) - } + self.truncate(new_cap); } let new_alloc = if self.header.cap == 0 { @@ -470,13 +486,19 @@ impl RawVector { } }; - self.header.cap = new_cap as u32; + // If the new capacity is lower than the length, we already dropped + // the elements after the new capacity, so make sure to write the + // length before propagating a potential allocation error. self.header.len = self.header.len.min(new_cap as u32); + // If allocation failed, we return an error here. let new_alloc = new_alloc?; + // The data and capacity, however must be only written once we know that + // the grow/shrink operation suceeded. let new_data_ptr = crate::raw::data_ptr::, T>(new_alloc.cast()); self.data = NonNull::new_unchecked(new_data_ptr); + self.header.cap = new_cap as u32; Ok(()) } @@ -522,12 +544,7 @@ impl RawVector { let old_len = self.len(); if old_len > new_cap { - let mut elt = self.data_ptr().add(new_cap); - let end = self.data_ptr().add(old_len); - while elt < end { - ptr::drop_in_place(elt); - elt = elt.add(1) - } + self.truncate(new_cap); } let old_cap = self.capacity(); @@ -1771,3 +1788,66 @@ fn reallocate_in() { assert_eq!(vec.capacity(), 0); assert_eq!(Rc::strong_count(&val), 1); } + +#[test] +fn extend_within_capacity() { + use std::sync::atomic::{AtomicUsize, Ordering}; + static COUNTER: AtomicUsize = AtomicUsize::new(0); + + struct Foo; + impl Drop for Foo { + fn drop(&mut self) { + COUNTER.fetch_add(1, Ordering::SeqCst); + } + } + + let values = [ + Foo, + Foo, + Foo, + Foo, + Foo, + Foo, + Foo, + Foo, + Foo, + Foo, + Foo, + Foo, + Foo, + Foo, + Foo, + Foo, + ]; + + let n = values.len(); + let cap; + + let mut vector = RawVector::new(); + unsafe { + vector.try_reserve(&Global, 4).unwrap(); + cap = vector.capacity(); + // What's mportant here is that we exercise the code path + // where we extend within a capcity that is inferior to the + // amount of elements that we need to push. So if the default + // capacity grows and this fails, just increase the number of + // elements in values. + assert!(cap < n); + + let mut iter = values.into_iter(); + + vector.extend_within_capacity(&mut iter); + + assert_eq!(COUNTER.load(Ordering::SeqCst), 0); + } + + assert_eq!(COUNTER.load(Ordering::SeqCst), n - cap); + + vector.clear(); + + assert_eq!(COUNTER.load(Ordering::SeqCst), n); + + unsafe { + vector.deallocate(&Global); + } +}