smallbox

Cargo.toml

@@ -4,7 +4,7 @@ version = "0.1.0"
 edition = "2024"
 
 [features]
-default = []
+default = ["alloc"]
 alloc = []
 std = []
 

src/lib.rs

@@ -1,11 +1,20 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
-#[cfg(feature = "alloc")]
+#[cfg(any(test, feature = "std", feature = "alloc"))]
 extern crate alloc;
 
-#[cfg(feature = "std")]
+#[cfg(any(test, feature = "std"))]
 extern crate std;
 
 pub mod cachepadded;
 pub mod drop_guard;
 pub mod ptr;
+#[cfg(feature = "alloc")]
+pub mod smallbox;
+
+pub const fn can_transmute<A, B>() -> bool {
+    use core::mem::{align_of, size_of};
+    // We can transmute `A` to `B` iff `A` and `B` have the same size and the
+    // alignment of `A` is greater than or equal to the alignment of `B`.
+    size_of::<A>() <= size_of::<B>() && align_of::<A>() >= align_of::<B>()
+}

src/smallbox.rs

@@ -0,0 +1,240 @@
+use core::{
+    borrow::{Borrow, BorrowMut},
+    cmp::Ordering,
+    mem::{self, ManuallyDrop, MaybeUninit},
+    ops::{Deref, DerefMut},
+};
+
+use alloc::boxed::Box;
+use core::fmt;
+
+/// A small box that can store a value inline if the size and alignment of T is
+/// less than or equal to the size and alignment of a boxed type. Typically this
+/// will be `sizeof::<usize>()` bytes, but might be larger if
+/// `sizeof::<Box<T>>()` is larger than that, like it is for dynamically sized
+/// types like `[T]` or `dyn Trait`.
+#[derive(Debug)]
+#[repr(transparent)]
+// We use a box here because a box can be unboxed, while a pointer cannot.
+pub struct SmallBox<T>(pub MaybeUninit<Box<T>>);
+
+impl<T: fmt::Display> fmt::Display for SmallBox<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        (**self).fmt(f)
+    }
+}
+
+impl<T: Ord> Ord for SmallBox<T> {
+    fn cmp(&self, other: &Self) -> Ordering {
+        self.as_ref().cmp(other.as_ref())
+    }
+}
+
+impl<T: PartialOrd> PartialOrd for SmallBox<T> {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        self.as_ref().partial_cmp(other.as_ref())
+    }
+}
+
+impl<T: Eq> Eq for SmallBox<T> {}
+
+impl<T: PartialEq> PartialEq for SmallBox<T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.as_ref().eq(other.as_ref())
+    }
+}
+
+impl<T: Default> Default for SmallBox<T> {
+    fn default() -> Self {
+        Self::new(Default::default())
+    }
+}
+
+impl<T: Clone> Clone for SmallBox<T> {
+    fn clone(&self) -> Self {
+        Self::new(self.as_ref().clone())
+    }
+}
+
+impl<T> Deref for SmallBox<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        self.as_ref()
+    }
+}
+
+impl<T> DerefMut for SmallBox<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.as_mut()
+    }
+}
+
+impl<T> AsRef<T> for SmallBox<T> {
+    fn as_ref(&self) -> &T {
+        Self::as_ref(self)
+    }
+}
+impl<T> AsMut<T> for SmallBox<T> {
+    fn as_mut(&mut self) -> &mut T {
+        Self::as_mut(self)
+    }
+}
+
+impl<T> Borrow<T> for SmallBox<T> {
+    fn borrow(&self) -> &T {
+        &**self
+    }
+}
+impl<T> BorrowMut<T> for SmallBox<T> {
+    fn borrow_mut(&mut self) -> &mut T {
+        &mut **self
+    }
+}
+
+impl<T> SmallBox<T> {
+    /// must only be called once. takes a reference so this can be called in
+    /// drop()
+    unsafe fn get_unchecked(&self, inline: bool) -> T {
+        if inline {
+            unsafe { mem::transmute_copy::<MaybeUninit<Box<T>>, T>(&self.0) }
+        } else {
+            unsafe { *self.0.assume_init_read() }
+        }
+    }
+
+    pub fn as_ref(&self) -> &T {
+        unsafe {
+            if Self::is_inline() {
+                mem::transmute::<&MaybeUninit<Box<T>>, &T>(&self.0)
+            } else {
+                self.0.assume_init_ref()
+            }
+        }
+    }
+
+    pub fn as_mut(&mut self) -> &mut T {
+        unsafe {
+            if Self::is_inline() {
+                mem::transmute::<&mut MaybeUninit<Box<T>>, &mut T>(&mut self.0)
+            } else {
+                self.0.assume_init_mut()
+            }
+        }
+    }
+
+    pub fn into_inner(self) -> T {
+        let this = ManuallyDrop::new(self);
+        let inline = Self::is_inline();
+
+        // SAFETY: inline is correctly calculated and this function
+        // consumes `self`
+        unsafe { this.get_unchecked(inline) }
+    }
+
+    #[inline(always)]
+    pub const fn is_inline() -> bool {
+        // the value can be stored inline iff the size of T is equal or
+        // smaller than the size of the boxed type and the alignment of the
+        // boxed type is an integer multiple of the alignment of T
+        crate::can_transmute::<Box<MaybeUninit<T>>, T>()
+    }
+
+    pub fn new(value: T) -> Self {
+        let inline = Self::is_inline();
+
+        if inline {
+            let mut this = MaybeUninit::new(Self(MaybeUninit::uninit()));
+            unsafe {
+                this.as_mut_ptr().cast::<T>().write(value);
+                this.assume_init()
+            }
+        } else {
+            Self(MaybeUninit::new(Box::new(value)))
+        }
+    }
+}
+
+impl<T> Drop for SmallBox<T> {
+    fn drop(&mut self) {
+        unsafe {
+            if Self::is_inline() {
+                mem::transmute::<_, &mut MaybeUninit<T>>(&mut self.0).assume_init_drop();
+            } else {
+                self.0.assume_init_drop();
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use core::sync::atomic::AtomicUsize;
+
+    use super::*;
+
+    #[test]
+    fn value_inline() {
+        assert!(SmallBox::<u32>::is_inline(), "u32 should be inline");
+        assert!(SmallBox::<u8>::is_inline(), "u8 should be inline");
+        assert!(
+            SmallBox::<Box<u32>>::is_inline(),
+            "Box<u32> should be inline"
+        );
+        assert!(
+            SmallBox::<[u32; 2]>::is_inline(),
+            "[u32; 2] should be inline"
+        );
+        assert!(
+            !SmallBox::<[u32; 3]>::is_inline(),
+            "[u32; 3] should not be inline"
+        );
+        assert!(SmallBox::<usize>::is_inline(), "usize should be inline");
+
+        #[repr(C, align(16))]
+        struct LargeType(u8);
+        assert!(
+            !SmallBox::<LargeType>::is_inline(),
+            "LargeType should not be inline"
+        );
+
+        #[repr(C, align(4))]
+        struct SmallType(u8);
+        assert!(
+            SmallBox::<SmallType>::is_inline(),
+            "SmallType should be inline"
+        );
+    }
+
+    #[test]
+    fn unbox_smallbox() {
+        let small_box = SmallBox::new(42u32);
+        assert_eq!(*small_box, 42);
+        let value: u32 = small_box.into_inner();
+        assert_eq!(value, 42);
+    }
+
+    #[test]
+    fn drop_smallbox() {
+        struct DropCounter<'a> {
+            count: &'a AtomicUsize,
+        }
+        impl<'a> Drop for DropCounter<'a> {
+            fn drop(&mut self) {
+                self.count
+                    .fetch_add(1, core::sync::atomic::Ordering::Relaxed);
+            }
+        }
+
+        let drop_count = AtomicUsize::new(0);
+        {
+            let _small_box = SmallBox::new(DropCounter { count: &drop_count });
+            assert_eq!(drop_count.load(core::sync::atomic::Ordering::Relaxed), 0);
+        }
+        assert_eq!(
+            drop_count.load(core::sync::atomic::Ordering::Relaxed),
+            1,
+            "DropCounter should have been dropped"
+        );
+    }
+}