A few fixes here and there.

src/drain.rs

@@ -104,17 +104,17 @@ impl<T> 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::<T>() == 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);

src/shared.rs

@@ -256,7 +256,7 @@ impl<T: Clone, R: RefCount, A: Allocator + Clone> RefCountedVector<T, R, A> {
         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<T, R: RefCount, A: Allocator> Drop for RefCountedVector<T, R, A> {
 }
 
 
-unsafe impl<T: Sync, A: Allocator + Send> Send for AtomicSharedVector<T, A> {}
+unsafe impl<T: Send + Sync, A: Allocator + Send> Send for AtomicSharedVector<T, A> {}
 
 unsafe impl<T: Send + Sync, A: Allocator + Sync> Sync for AtomicSharedVector<T, A> {}
 

src/vector.rs

@@ -378,18 +378,21 @@ impl<T> RawVector<T> {
 
     unsafe fn extend_within_capacity(&mut self, iter: &mut impl Iterator<Item = T>) {
         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<T> RawVector<T> {
         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<A: Allocator>(&mut self, allocator: &A, new_cap: usize) -> Result<(), AllocError> {
         type R = DefaultRefCount;
@@ -447,12 +468,7 @@ impl<T> RawVector<T> {
 
         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<T> RawVector<T> {
             }
         };
 
-        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::<Header<R, A>, 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<T> RawVector<T> {
 
         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);
+    }
+}