smallbox

smallbox.rs

// rustc:

use core::marker::PhantomData;
use core::mem::{ManuallyDrop, MaybeUninit};
use core::ops::{Deref, DerefMut};
use core::ptr::{addr_of, addr_of_mut};

#[derive(Copy, Clone, Debug)]
enum Kind {
    Inline,
    InlineDst(*const ()),
    Boxed,
    BoxedDst(*const ()),
}

#[derive(Debug)]
struct SmallBox<T: ?Sized, S: Storage> {
    kind: Kind,
    repr: Repr<S>,
    _marker: PhantomData<Box<T>>,
}

#[derive(Copy, Clone)]
#[repr(C)]
union Repr<S> {
    boxed: *mut (),
    storage: ManuallyDrop<MaybeUninit<S>>,
}

impl<S> core::fmt::Debug for Repr<S> {
    fn fmt(&self, fmt: &mut core::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
        fmt.debug_struct("Repr").finish_non_exhaustive()
    }
}

#[repr(C)]
union StorageOrVal<T, S> {
    storage: ManuallyDrop<MaybeUninit<S>>,
    val: ManuallyDrop<T>,
}

impl<T: ?Sized, S: Storage> SmallBox<T, S> {
    unsafe fn new_dst<D>(dst: *const T, sized: SmallBox<D, S>) -> Self {
        const {
            assert!(
                core::mem::align_of::<*const T>() == core::mem::align_of::<*const ()>(),
                "The destination must be a DST."
            );
        }

        let sized = ManuallyDrop::new(sized);

        let is_dst =
            core::mem::size_of::<*const T>() == core::mem::size_of::<(*const (), *const ())>();
        let metadata: Option<*const ()> = if is_dst {
            Some(unsafe { PtrRepr { const_ptr: dst }.components.metadata })
        } else {
            None
        };

        let kind = match sized.kind {
            Kind::Inline | Kind::InlineDst(_) => {
                if let Some(metadata) = metadata {
                    Kind::InlineDst(metadata)
                } else {
                    Kind::Inline
                }
            }
            Kind::Boxed | Kind::BoxedDst(_) => {
                if let Some(metadata) = metadata {
                    Kind::BoxedDst(metadata)
                } else {
                    Kind::Boxed
                }
            }
        };

        Self {
            kind,
            repr: sized.repr,
            _marker: PhantomData,
        }
    }

    fn as_ptr(&self) -> *const T {
        self.as_ptr_inner().0
    }

    fn as_ptr_inner(&self) -> (*const T, bool) {
        match self.kind {
            Kind::Inline => unsafe {
                (
                    PtrRepr {
                        components: PtrComponents {
                            ptr: addr_of!(self.repr.storage) as *const (),
                            metadata: (),
                        },
                    }
                    .const_ptr,
                    false,
                )
            },
            Kind::InlineDst(metadata) => unsafe {
                (
                    PtrRepr {
                        components: PtrComponents {
                            ptr: addr_of!(self.repr.storage) as *const (),
                            metadata,
                        },
                    }
                    .const_ptr,
                    false,
                )
            },
            Kind::Boxed => unsafe {
                (
                    PtrRepr {
                        components: PtrComponents {
                            ptr: self.repr.boxed,
                            metadata: (),
                        },
                    }
                    .const_ptr,
                    true,
                )
            },
            Kind::BoxedDst(metadata) => unsafe {
                (
                    PtrRepr {
                        components: PtrComponents {
                            ptr: self.repr.boxed,
                            metadata,
                        },
                    }
                    .const_ptr,
                    true,
                )
            },
        }
    }

    fn as_mut_ptr(&mut self) -> *mut T {
        self.as_mut_ptr_inner().0
    }

    fn as_mut_ptr_inner(&mut self) -> (*mut T, bool) {
        match self.kind {
            Kind::Inline => unsafe {
                (
                    PtrRepr {
                        components: PtrComponents {
                            ptr: addr_of_mut!(self.repr.storage) as *mut (),
                            metadata: (),
                        },
                    }
                    .mut_ptr,
                    false,
                )
            },
            Kind::InlineDst(metadata) => unsafe {
                (
                    PtrRepr {
                        components: PtrComponents {
                            ptr: addr_of_mut!(self.repr.storage) as *mut (),
                            metadata,
                        },
                    }
                    .mut_ptr,
                    false,
                )
            },
            Kind::Boxed => unsafe {
                (
                    PtrRepr {
                        components: PtrComponents {
                            ptr: self.repr.boxed,
                            metadata: (),
                        },
                    }
                    .mut_ptr,
                    true,
                )
            },
            Kind::BoxedDst(metadata) => unsafe {
                (
                    PtrRepr {
                        components: PtrComponents {
                            ptr: self.repr.boxed,
                            metadata,
                        },
                    }
                    .mut_ptr,
                    true,
                )
            },
        }
    }

    fn as_ref(&self) -> &T {
        unsafe { &*self.as_ptr() }
    }

    fn as_mut_ref(&mut self) -> &mut T {
        unsafe { &mut *self.as_mut_ptr() }
    }
}

impl<T, S: Storage> SmallBox<T, S> {
    const IS_INLINABLE: bool = may_store::<S, T>();

    fn new(val: T) -> Self {
        if Self::IS_INLINABLE {
            unsafe { Self::new_inline(val) }
        } else {
            unsafe { Self::new_boxed(val) }
        }
    }

    unsafe fn new_inline(val: T) -> Self {
        Self {
            kind: Kind::Inline,
            repr: Repr {
                storage: unsafe {
                    StorageOrVal {
                        val: ManuallyDrop::new(val),
                    }
                    .storage
                },
            },
            _marker: PhantomData,
        }
    }

    unsafe fn new_boxed(val: T) -> Self {
        Self {
            kind: Kind::Boxed,
            repr: Repr {
                boxed: Box::into_raw(Box::new(val)).cast::<()>(),
            },
            _marker: PhantomData,
        }
    }
}

impl<T: ?Sized, S: Storage> Deref for SmallBox<T, S> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.as_ref()
    }
}

impl<T: ?Sized, S: Storage> DerefMut for SmallBox<T, S> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.as_mut_ref()
    }
}

impl<T: ?Sized, S: Storage> Drop for SmallBox<T, S> {
    fn drop(&mut self) {
        let (ptr, is_boxed) = self.as_mut_ptr_inner();

        unsafe {
            if is_boxed {
                drop(Box::from_raw(ptr));
            } else {
                core::ptr::drop_in_place(self.as_mut_ptr());
            }
        }
    }
}

mod private {
    pub trait Sealed {}
}

trait Storage: Clone + Copy + private::Sealed {
    const ALIGN: usize;
    const SIZE: usize;
}

const fn may_store<S: Storage, T>() -> bool {
    core::mem::size_of::<T>() <= S::SIZE && core::mem::align_of::<T>() <= S::ALIGN
}

impl<const N: usize, T: Copy> Storage for [T; N] {
    const ALIGN: usize = core::mem::align_of::<T>();
    const SIZE: usize = core::mem::size_of::<Self>();
}

impl<const N: usize, T> private::Sealed for [T; N] {}

#[repr(C)]
union PtrRepr<T: ?Sized, M: Copy> {
    const_ptr: *const T,
    mut_ptr: *mut T,
    components: PtrComponents<M>,
}

#[derive(Clone, Copy)]
#[repr(C)]
struct PtrComponents<M: Copy> {
    ptr: *const (),
    metadata: M,
}

macro_rules! smallbox_cast {
    ($b:expr, $t:ty) => {{
        let moved = $b;
        unsafe { SmallBox::<$t, _>::new_dst(SmallBox::as_ref(&moved) as &$t as *const _, moved) }
    }};
}

trait Inc: core::fmt::Debug {
    fn inc(&mut self);

    fn as_debug(&self) -> &dyn core::fmt::Debug;
}

impl Inc for usize {
    fn inc(&mut self) {
        *self += 1;
    }

    fn as_debug(&self) -> &dyn core::fmt::Debug {
        self as _
    }
}

fn main() {
    let m: SmallBox<usize, [usize; 2]> = SmallBox::new(1);

    let mut m = smallbox_cast!(m, usize);
    m.inc();

    let mut m = smallbox_cast!(m, dyn Inc);
    m.inc();

    println!("{:?}", m.as_debug());
    m.inc();
    println!("{:?}", m.as_debug());
    println!("{:?}", m);
}