diff --git a/src/smallbox.rs b/src/smallbox.rs index 74e34fe..1843da7 100644 --- a/src/smallbox.rs +++ b/src/smallbox.rs @@ -1,4 +1,3 @@ -use core::any::Any; use core::cell::UnsafeCell; use core::cmp::Ordering; use core::fmt; @@ -16,12 +15,26 @@ use core::ops; use core::ops::CoerceUnsized; use core::pin::Pin; use core::ptr; +use core::ptr::NonNull; +use core::{any::Any, hint::unreachable_unchecked}; use ::alloc::alloc; use ::alloc::alloc::Layout; +use ::alloc::alloc::handle_alloc_error; use crate::sptr; +/// A sentinel pointer that signals that the value is stored on the stack +/// +/// It is never supposed to be dereferenced +const INLINE_SENTINEL: *mut u8 = 0x1 as *mut u8; + +/// Minimum alignment for allocations +/// +/// Forcing a minimum alignment prevents the allocator +/// from returning a pointer with the same address as `INLINE_SENTINEL` +const MIN_ALIGNMENT: usize = 2; + #[cfg(feature = "coerce")] impl, U: ?Sized, Space> CoerceUnsized> for SmallBox @@ -71,7 +84,8 @@ macro_rules! smallbox { /// An optimized box that store value on stack or on heap depending on its size pub struct SmallBox { space: MaybeUninit>, - ptr: *const T, + // NonNull enables Null Pointer Optimization + ptr: NonNull, _phantom: PhantomData, } @@ -100,14 +114,18 @@ impl SmallBox { /// ``` #[inline(always)] pub fn new(val: T) -> SmallBox - where T: Sized { + where + T: Sized, + { smallbox!(val) } #[doc(hidden)] #[inline] pub unsafe fn new_unchecked(val: U, ptr: *const T) -> SmallBox - where U: Sized { + where + U: Sized, + { let val = ManuallyDrop::new(val); Self::new_copy(&val, ptr) } @@ -161,31 +179,42 @@ impl SmallBox { /// ``` #[inline] pub fn is_heap(&self) -> bool { - !self.ptr.is_null() + self.ptr.as_ptr() as *mut u8 != INLINE_SENTINEL } unsafe fn new_copy(val: &U, metadata_ptr: *const T) -> SmallBox - where U: ?Sized { + where + U: ?Sized, + { let size = mem::size_of_val::(val); let align = mem::align_of_val::(val); let mut space = MaybeUninit::>::uninit(); let (ptr_this, val_dst): (*mut u8, *mut u8) = if size == 0 { - (ptr::null_mut(), sptr::without_provenance_mut(align)) + (INLINE_SENTINEL, sptr::without_provenance_mut(align)) } else if size > mem::size_of::() || align > mem::align_of::() { // Heap - let layout = Layout::for_value::(val); + // Safety: MIN_ALIGNMENT is 2, aligning to 2 should not create an invalid layout + let layout = Layout::for_value::(val) + .align_to(MIN_ALIGNMENT) + .unwrap_or_else(|_| unreachable_unchecked()); let heap_ptr = alloc::alloc(layout); + if heap_ptr.is_null() { + handle_alloc_error(layout) + } + (heap_ptr, heap_ptr) } else { // Stack - (ptr::null_mut(), space.as_mut_ptr().cast()) + (INLINE_SENTINEL, space.as_mut_ptr().cast()) }; // `self.ptr` always holds the metadata, even if stack allocated let ptr = sptr::with_metadata_of_mut(ptr_this, metadata_ptr); + // Safety: is either an INLINE_SENTINEL or is returned from the allocator and is checked for null + let ptr = NonNull::new_unchecked(ptr); ptr::copy_nonoverlapping(sptr::from_ref(val).cast(), val_dst, size); @@ -222,18 +251,18 @@ impl SmallBox { #[inline] unsafe fn as_ptr(&self) -> *const T { if self.is_heap() { - self.ptr + self.ptr.as_ptr() } else { - sptr::with_metadata_of(self.space.as_ptr(), self.ptr) + sptr::with_metadata_of(self.space.as_ptr(), self.ptr.as_ptr()) } } #[inline] unsafe fn as_mut_ptr(&mut self) -> *mut T { if self.is_heap() { - self.ptr as *mut T + self.ptr.as_ptr() } else { - sptr::with_metadata_of_mut(self.space.as_mut_ptr(), self.ptr) + sptr::with_metadata_of_mut(self.space.as_mut_ptr(), self.ptr.as_ptr()) } } @@ -254,15 +283,22 @@ impl SmallBox { /// ``` #[inline] pub fn into_inner(self) -> T - where T: Sized { + where + T: Sized, + { let this = ManuallyDrop::new(self); let ret_val: T = unsafe { this.as_ptr().read() }; // Just drops the heap without dropping the boxed value if this.is_heap() { - let layout = Layout::new::(); + // Safety: MIN_ALIGNMENT is 2, aligning to 2 should not create an invalid layout + let layout = unsafe { + Layout::new::() + .align_to(MIN_ALIGNMENT) + .unwrap_or_else(|_| unreachable_unchecked()) + }; unsafe { - alloc::dealloc(this.ptr as *const u8 as *mut u8, layout); + alloc::dealloc(this.ptr.as_ptr() as *const u8 as *mut u8, layout); } } @@ -363,17 +399,21 @@ impl ops::DerefMut for SmallBox { impl ops::Drop for SmallBox { fn drop(&mut self) { unsafe { - let layout = Layout::for_value::(&*self); + let layout = Layout::for_value::(&*self) + .align_to(MIN_ALIGNMENT) + .unwrap_or_else(|_| unreachable_unchecked()); + ptr::drop_in_place::(&mut **self); if self.is_heap() { - alloc::dealloc(self.ptr as *const u8 as *mut u8, layout); + alloc::dealloc(self.ptr.as_ptr() as *const u8 as *mut u8, layout); } } } } impl Clone for SmallBox -where T: Sized +where + T: Sized, { fn clone(&self) -> Self { let val: &T = self; @@ -465,8 +505,8 @@ unsafe impl Sync for SmallBox {} #[cfg(test)] mod tests { - use core::any::Any; use core::ptr::addr_of; + use core::{any::Any, mem}; use ::alloc::boxed::Box; use ::alloc::vec; @@ -702,4 +742,12 @@ mod tests { val } } + + #[test] + fn test_null_ptr_optimization() { + assert_eq!( + mem::size_of::>(), + mem::size_of::>>() + ); + } }