Fixed size box implementation?

I want a store a dyn Foo directly inside a struct (to improve cache locality compared to boxing). I know that since trait objects don't have a known size that my struct will only support some fixed upper limit. I'm imagining a type that looks something like FixedSizeBox<dyn Foo, 24> that could store any implementer of Foo that is 24 bytes or smaller. Is there a crate that does something like this already? I found dynstack and heapless but neither are really targeting this case.

I don't think this is possible to soundly implement in stable Rust, since we cannot extract the trait metadata from incoming values. But in nightly Rust, we can implement a FixedSizeBox with the ptr_metadata feature, and we can convert an owned object T into a FixedSizeBox with the unsize feature (Rust Playground):

#![feature(ptr_metadata)]
#![feature(unsize)]

use std::{
    alloc, fmt,
    marker::Unsize,
    mem::{self, ManuallyDrop, MaybeUninit},
    ops::{Deref, DerefMut},
    ptr::{self, DynMetadata, Pointee},
};

#[repr(C, align(16))]
pub struct FixedSizeBox<T: ?Sized, const N: usize>
where
    T: Pointee<Metadata = DynMetadata<T>>,
{
    data: [MaybeUninit<u8>; N],
    metadata: DynMetadata<T>,
}

impl<T: ?Sized, const N: usize> FixedSizeBox<T, N>
where
    T: Pointee<Metadata = DynMetadata<T>>,
{
    pub fn new<U: Unsize<T>>(value: U) -> Self {
        assert!(mem::size_of::<U>() <= N, "object too large");
        assert!(mem::align_of::<U>() <= 16, "incompatible alignment");
        let metadata = ptr::metadata(&value as &T);
        let mut data: [MaybeUninit<u8>; N] = unsafe { MaybeUninit::uninit().assume_init() };
        unsafe {
            data.as_mut_ptr().cast::<U>().write(value);
        }
        FixedSizeBox { data, metadata }
    }

    pub fn from_box(b: Box<T>) -> Self {
        let metadata = ptr::metadata(&*b);
        assert!(metadata.size_of() <= N, "object too large");
        assert!(metadata.align_of() <= 16, "incompatible alignment");
        let src = Box::into_raw(b).cast::<u8>();
        let mut data: [MaybeUninit<u8>; N] = unsafe { MaybeUninit::uninit().assume_init() };
        unsafe {
            ptr::copy_nonoverlapping(src, data.as_mut_ptr().cast::<u8>(), metadata.size_of());
            alloc::dealloc(src, metadata.layout());
        }
        FixedSizeBox { data, metadata }
    }

    pub fn into_box(this: Self) -> Box<T> {
        let this = ManuallyDrop::new(this);
        unsafe {
            let dest = alloc::alloc(this.metadata.layout());
            ptr::copy_nonoverlapping(
                this.data.as_ptr().cast::<u8>(),
                dest,
                this.metadata.size_of(),
            );
            Box::from_raw(ptr::from_raw_parts_mut(dest.cast(), this.metadata))
        }
    }
}

impl<T: ?Sized, const N: usize> Drop for FixedSizeBox<T, N>
where
    T: Pointee<Metadata = DynMetadata<T>>,
{
    fn drop(&mut self) {
        unsafe { ptr::drop_in_place(&mut **self) }
    }
}

impl<T: ?Sized, const N: usize> Deref for FixedSizeBox<T, N>
where
    T: Pointee<Metadata = DynMetadata<T>>,
{
    type Target = T;

    fn deref(&self) -> &Self::Target {
        let ptr = ptr::from_raw_parts(self.data.as_ptr().cast(), self.metadata);
        unsafe { &*ptr }
    }
}

impl<T: ?Sized, const N: usize> DerefMut for FixedSizeBox<T, N>
where
    T: Pointee<Metadata = DynMetadata<T>>,
{
    fn deref_mut(&mut self) -> &mut Self::Target {
        let ptr = ptr::from_raw_parts_mut(self.data.as_mut_ptr().cast(), self.metadata);
        unsafe { &mut *ptr }
    }
}

impl<T: ?Sized, const N: usize> fmt::Debug for FixedSizeBox<T, N>
where
    T: Pointee<Metadata = DynMetadata<T>> + fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        (**self).fmt(f)
    }
}

fn main() {
    let s: FixedSizeBox<dyn fmt::Debug, 24> = FixedSizeBox::new("Hello, world!");
    println!("{:?}", s);
    let s: Box<dyn fmt::Debug> = FixedSizeBox::into_box(s);
    println!("{:?}", s);
}

I'm pretty sure this is sound as written, but I haven't inspected it closely enough to give any absolute guarantees. It has the caveat that the stored object has a fixed maximum alignment; this could be fixed with a bit more logic to store the object at an offset in the data array.

4 Likes

There are alternatives around a boxed dyn Trait. You can use generics :

struct Holds<F: Foo> {
    // other fields,
    a_foo: F,
}

You can make an enum of all possible <24 byte variants:

enum SmallFooType {
    FooA(FooA),
    FooB(FooB),
    BigFoo(Box<dyn BigFoo>), // Boxing the big ones may be worthwhile
}

There's even a crate called enum_dispatch that helps handle an enum of objects that implement a trait.

It might also be worth pointing out that because dyn Trait is a wall of pointer indirection (to the base struct and to the vtable of methods of the trait), compiler optimizations such as inline calling of small methods is lost, so might encourage generics or an enum.

And, while doing a quick search for the Rust book bit on object safety of traits, I came across this nice summary of the different options: 3 Things to Try When You Can't Make a Trait Object

1 Like

If you could require an additional bound for something like Into<FixedSizeBox<N>>, that doesn't seem like it would be too hard to blanket implement, but it would be a pain to thread the requirements everywhere.

Thanks this is exactly what I was looking for. Still surprised nobody has put it up as a crate yet.