fix: drop, clone, default and f calls when allocating ZSTs #6

bluurryy · Apr 7, 2024 · 5e3ed18 · 5e3ed18
1 parent db136de
commit 5e3ed18
Show file tree

Hide file tree

Showing 6 changed files with 252 additions and 29 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## Unreleased
+- **fixed:** ZST allocation with respect to `drop`, `clone` and `default` calls
+- **fixed:** `alloc_with` and `alloc_slice_fill_with` not calling `f` for ZSTs
+
 ## 0.1.5 (2024-04-05)
 - **added:** `BumpVec::into_fixed_vec` and `FixedBumpVec::into_vec`
 - **added:** fallible `FixedBumpVec` api

diff --git a/src/bump_box.rs b/src/bump_box.rs
@@ -139,8 +139,9 @@ pub struct BumpBox<'a, T: ?Sized> {
 impl<'a, T> BumpBox<'a, T> {
     #[must_use]
     #[inline(always)]
-    pub(crate) fn zst() -> Self {
+    pub(crate) fn zst(value: T) -> Self {
         assert!(T::IS_ZST);
+        mem::forget(value);
         Self {
             ptr: NonNull::dangling(),
             marker: PhantomData,
@@ -323,6 +324,16 @@ impl<'a, T: Sized> BumpBox<'a, MaybeUninit<T>> {
 }
 
 impl<'a, T: Sized> BumpBox<'a, [MaybeUninit<T>]> {
+    #[must_use]
+    #[inline(always)]
+    pub(crate) fn uninit_zst_slice(len: usize) -> Self {
+        assert!(T::IS_ZST);
+        Self {
+            ptr: nonnull::slice_from_raw_parts(NonNull::dangling(), len),
+            marker: PhantomData,
+        }
+    }
+
     /// Initializes `self` by filling it with elements by cloning `value`.
     ///
     /// # Examples
@@ -477,7 +488,48 @@ impl<'a, T> BumpBox<'a, [T]> {
 
     #[must_use]
     #[inline(always)]
-    pub(crate) fn zst_slice(len: usize) -> Self {
+    pub(crate) fn zst_slice_clone(slice: &[T]) -> Self
+    where
+        T: Clone,
+    {
+        assert!(T::IS_ZST);
+        BumpBox::uninit_zst_slice(slice.len()).init_clone(slice)
+    }
+
+    #[must_use]
+    #[inline(always)]
+    pub(crate) fn zst_slice_fill(len: usize, value: T) -> Self
+    where
+        T: Clone,
+    {
+        assert!(T::IS_ZST);
+        if len == 0 {
+            drop(value);
+            BumpBox::EMPTY
+        } else {
+            for _ in 1..len {
+                mem::forget(value.clone());
+            }
+
+            mem::forget(value);
+            unsafe { BumpBox::zst_slice_from_len(len) }
+        }
+    }
+
+    #[must_use]
+    #[inline(always)]
+    pub(crate) fn zst_slice_fill_with(len: usize, mut f: impl FnMut() -> T) -> Self {
+        assert!(T::IS_ZST);
+        for _ in 0..len {
+            mem::forget(f());
+        }
+        unsafe { BumpBox::zst_slice_from_len(len) }
+    }
+
+    /// Creates `T` values from nothing!
+    #[must_use]
+    #[inline(always)]
+    unsafe fn zst_slice_from_len(len: usize) -> Self {
         assert!(T::IS_ZST);
         Self {
             ptr: nonnull::slice_from_raw_parts(NonNull::dangling(), len),
@@ -1036,7 +1088,7 @@ impl<'a, T> BumpBox<'a, [T]> {
             let _ = self.into_raw();
             let _ = other.into_raw();
 
-            Self::zst_slice(len)
+            unsafe { Self::zst_slice_from_len(len) }
         } else {
             if self.as_ptr_range().end != other.as_ptr() {
                 assert_failed();

diff --git a/src/lib.rs b/src/lib.rs
@@ -1222,7 +1222,7 @@ define_alloc_methods! {
     for pub fn try_alloc
     fn generic_alloc<{T}>(&self, value: T) -> BumpBox<T> | BumpBox<'a, T> {
         if T::IS_ZST {
-            return Ok(BumpBox::zst());
+            return Ok(BumpBox::zst(value));
         }
 
         self.generic_alloc_with(|| value)
@@ -1235,7 +1235,8 @@ define_alloc_methods! {
     for pub fn try_alloc_with
     fn generic_alloc_with<{T}>(&self, f: impl FnOnce() -> T) -> BumpBox<T> | BumpBox<'a, T> {
         if T::IS_ZST {
-            return Ok(BumpBox::zst());
+            let value = f();
+            return Ok(BumpBox::zst(value));
         }
 
         let ptr = self.do_alloc_sized::<B, T>()?;
@@ -1263,7 +1264,7 @@ define_alloc_methods! {
     fn generic_alloc_try_with<{T, E}>(&self, f: impl FnOnce() -> Result<T, E>) -> Result<BumpBox<T>, E> | Result<BumpBox<'a, T>, E> {
         if T::IS_ZST {
             return match f() {
-                Ok(_) => Ok(Ok(BumpBox::zst())),
+                Ok(value) => Ok(Ok(BumpBox::zst(value))),
                 Err(error) => Ok(Err(error)),
             }
         }
@@ -1308,7 +1309,7 @@ define_alloc_methods! {
         slice: &[T],
     ) -> BumpBox<[T]> | BumpBox<'a, [T]> {
         if T::IS_ZST {
-            return Ok(BumpBox::zst_slice(slice.len()));
+            return Ok(BumpBox::zst_slice_clone(slice));
         }
 
         let len = slice.len();
@@ -1330,7 +1331,7 @@ define_alloc_methods! {
         slice: &[T],
     ) -> BumpBox<[T]> | BumpBox<'a, [T]> {
         if T::IS_ZST {
-            return Ok(BumpBox::zst_slice(slice.len()));
+            return Ok(BumpBox::zst_slice_clone(slice));
         }
 
         Ok(self.generic_alloc_uninit_slice_for(slice)?.init_clone(slice))
@@ -1342,7 +1343,7 @@ define_alloc_methods! {
     for pub fn try_alloc_slice_fill
     fn generic_alloc_slice_fill<{T: Clone}>(&self, len: usize, value: T) -> BumpBox<[T]> | BumpBox<'a, [T]> {
         if T::IS_ZST {
-            return Ok(BumpBox::zst_slice(len));
+            return Ok(BumpBox::zst_slice_fill(len, value));
         }
 
         Ok(self.generic_alloc_uninit_slice(len)?.init_fill(value))
@@ -1359,7 +1360,7 @@ define_alloc_methods! {
     for pub fn try_alloc_slice_fill_with
     fn generic_alloc_slice_fill_with<{T}>(&self, len: usize, f: impl FnMut() -> T) -> BumpBox<[T]> | BumpBox<'a, [T]> {
         if T::IS_ZST {
-            return Ok(BumpBox::zst_slice(len));
+            return Ok(BumpBox::zst_slice_fill_with(len, f));
         }
 
         Ok(self.generic_alloc_uninit_slice(len)?.init_fill_with(f))
@@ -1461,10 +1462,6 @@ define_alloc_methods! {
     /// For better performance prefer [`try_alloc_iter_exact`](Bump::try_alloc_iter_exact) or [`try_alloc_iter_mut(_rev)`](Bump::try_alloc_iter_mut).
     for pub fn try_alloc_iter
     fn generic_alloc_iter<{T}>(&self, iter: impl IntoIterator<Item = T>) -> BumpBox<[T]> | BumpBox<'a, [T]> {
-        if T::IS_ZST {
-            return Ok(BumpBox::zst_slice(iter.into_iter().count()));
-        }
-
         let iter = iter.into_iter();
         let capacity = iter.size_hint().0;
 
@@ -1495,10 +1492,6 @@ define_alloc_methods! {
     where {
         I: ExactSizeIterator<Item = T>
     } in {
-        if T::IS_ZST {
-            return Ok(BumpBox::zst_slice(iter.into_iter().count()));
-        }
-
         let mut iter = iter.into_iter();
         let len = iter.len();
 
@@ -1536,10 +1529,6 @@ define_alloc_methods! {
     /// When bumping downwards, prefer [`try_alloc_iter_mut_rev`](Bump::try_alloc_iter_mut_rev) or [`try_alloc_iter_exact`](Bump::try_alloc_iter_exact) as in this case this function incurs an additional copy of the slice internally.
     for pub fn try_alloc_iter_mut
     fn generic_alloc_iter_mut<{T}>(&mut self, iter: impl IntoIterator<Item = T>) -> BumpBox<[T]> | BumpBox<'a, [T]> {
-        if T::IS_ZST {
-            return Ok(BumpBox::zst_slice(iter.into_iter().count()));
-        }
-
         let iter = iter.into_iter();
         let capacity = iter.size_hint().0;
 
@@ -1569,10 +1558,6 @@ define_alloc_methods! {
     /// When bumping upwards, prefer [`try_alloc_iter_mut`](Bump::try_alloc_iter) or [`try_alloc_iter_exact`](Bump::try_alloc_iter_exact) as in this case this function incurs an additional copy of the slice internally.
     for pub fn try_alloc_iter_mut_rev
     fn generic_alloc_iter_mut_rev<{T}>(&mut self, iter: impl IntoIterator<Item = T>) -> BumpBox<[T]> | BumpBox<'a, [T]> {
-        if T::IS_ZST {
-            return Ok(BumpBox::zst_slice(iter.into_iter().count()));
-        }
-
         let iter = iter.into_iter();
         let capacity = iter.size_hint().0;
 
@@ -1618,7 +1603,7 @@ define_alloc_methods! {
     for pub fn try_alloc_uninit
     fn generic_alloc_uninit<{T}>(&self) -> BumpBox<MaybeUninit<T>> | BumpBox<'a, MaybeUninit<T>> {
         if T::IS_ZST {
-            return Ok(BumpBox::zst());
+            return Ok(BumpBox::zst(MaybeUninit::uninit()));
         }
 
         let ptr = self.do_alloc_sized::<B, T>()?.cast::<MaybeUninit<T>>();
@@ -1661,7 +1646,7 @@ define_alloc_methods! {
     for pub fn try_alloc_uninit_slice
     fn generic_alloc_uninit_slice<{T}>(&self, len: usize) -> BumpBox<[MaybeUninit<T>]> | BumpBox<'a, [MaybeUninit<T>]> {
         if T::IS_ZST {
-            return Ok(BumpBox::zst_slice(len));
+            return Ok(BumpBox::uninit_zst_slice(len));
         }
 
         let ptr = self.do_alloc_slice::<B, T>(len)?.cast::<MaybeUninit<T>>();
@@ -1683,7 +1668,7 @@ define_alloc_methods! {
     for pub fn try_alloc_uninit_slice_for
     fn generic_alloc_uninit_slice_for<{T}>(&self, slice: &[T]) -> BumpBox<[MaybeUninit<T>]> | BumpBox<'a, [MaybeUninit<T>]> {
         if T::IS_ZST {
-            return Ok(BumpBox::zst_slice(slice.len()));
+            return Ok(BumpBox::uninit_zst_slice(slice.len()));
         }
 
         let ptr = self.do_alloc_slice_for::<B, T>(slice)?.cast::<MaybeUninit<T>>();

diff --git a/src/mut_bump_vec.rs b/src/mut_bump_vec.rs
@@ -1025,6 +1025,10 @@ where
     fn into_slice_ptr(self) -> NonNull<[T]> {
         let mut this = ManuallyDrop::new(self);
 
+        if T::IS_ZST {
+            return nonnull::slice_from_raw_parts(NonNull::dangling(), this.len());
+        }
+
         if this.capacity() == 0 {
             // We didn't touch the bump, so no need to do anything.
             debug_assert_eq!(this.as_non_null_ptr(), NonNull::<T>::dangling());

diff --git a/src/mut_bump_vec_rev.rs b/src/mut_bump_vec_rev.rs
@@ -1101,6 +1101,10 @@ where
     fn into_slice_ptr(self) -> NonNull<[T]> {
         let this = ManuallyDrop::new(self);
 
+        if T::IS_ZST {
+            return nonnull::slice_from_raw_parts(NonNull::dangling(), this.len());
+        }
+
         if this.cap == 0 {
             // We didn't touch the bump, so no need to do anything.
             debug_assert_eq!(this.end, NonNull::<T>::dangling());