From 2fabc61378a9cc1f5f47cf38ae98245990af9603 Mon Sep 17 00:00:00 2001 From: Janis Date: Wed, 2 Jul 2025 18:04:21 +0200 Subject: [PATCH] smallbox --- Cargo.toml | 2 +- src/lib.rs | 13 ++- src/smallbox.rs | 240 ++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 252 insertions(+), 3 deletions(-) create mode 100644 src/smallbox.rs diff --git a/Cargo.toml b/Cargo.toml index 643d3cb..3c15d50 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,7 +4,7 @@ version = "0.1.0" edition = "2024" [features] -default = [] +default = ["alloc"] alloc = [] std = [] diff --git a/src/lib.rs b/src/lib.rs index c20f881..4513760 100644 --- a/src/lib.rs +++ b/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() -> 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::() <= size_of::() && align_of::() >= align_of::() +} diff --git a/src/smallbox.rs b/src/smallbox.rs new file mode 100644 index 0000000..8a6e678 --- /dev/null +++ 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::()` bytes, but might be larger if +/// `sizeof::>()` 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(pub MaybeUninit>); + +impl fmt::Display for SmallBox { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + (**self).fmt(f) + } +} + +impl Ord for SmallBox { + fn cmp(&self, other: &Self) -> Ordering { + self.as_ref().cmp(other.as_ref()) + } +} + +impl PartialOrd for SmallBox { + fn partial_cmp(&self, other: &Self) -> Option { + self.as_ref().partial_cmp(other.as_ref()) + } +} + +impl Eq for SmallBox {} + +impl PartialEq for SmallBox { + fn eq(&self, other: &Self) -> bool { + self.as_ref().eq(other.as_ref()) + } +} + +impl Default for SmallBox { + fn default() -> Self { + Self::new(Default::default()) + } +} + +impl Clone for SmallBox { + fn clone(&self) -> Self { + Self::new(self.as_ref().clone()) + } +} + +impl Deref for SmallBox { + type Target = T; + + fn deref(&self) -> &Self::Target { + self.as_ref() + } +} + +impl DerefMut for SmallBox { + fn deref_mut(&mut self) -> &mut Self::Target { + self.as_mut() + } +} + +impl AsRef for SmallBox { + fn as_ref(&self) -> &T { + Self::as_ref(self) + } +} +impl AsMut for SmallBox { + fn as_mut(&mut self) -> &mut T { + Self::as_mut(self) + } +} + +impl Borrow for SmallBox { + fn borrow(&self) -> &T { + &**self + } +} +impl BorrowMut for SmallBox { + fn borrow_mut(&mut self) -> &mut T { + &mut **self + } +} + +impl SmallBox { + /// 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::>, T>(&self.0) } + } else { + unsafe { *self.0.assume_init_read() } + } + } + + pub fn as_ref(&self) -> &T { + unsafe { + if Self::is_inline() { + mem::transmute::<&MaybeUninit>, &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>, &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::>, 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::().write(value); + this.assume_init() + } + } else { + Self(MaybeUninit::new(Box::new(value))) + } + } +} + +impl Drop for SmallBox { + fn drop(&mut self) { + unsafe { + if Self::is_inline() { + mem::transmute::<_, &mut MaybeUninit>(&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::::is_inline(), "u32 should be inline"); + assert!(SmallBox::::is_inline(), "u8 should be inline"); + assert!( + SmallBox::>::is_inline(), + "Box 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::::is_inline(), "usize should be inline"); + + #[repr(C, align(16))] + struct LargeType(u8); + assert!( + !SmallBox::::is_inline(), + "LargeType should not be inline" + ); + + #[repr(C, align(4))] + struct SmallType(u8); + assert!( + SmallBox::::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" + ); + } +}