Type-setting builder pattern for composition of traits/generics into single Box

code review
1

Suppose you have some objects A/B/C/D that all implement SomeTrait, as well as some more objects I/J/K/L/M/N that implement OtherTrait, and perhaps even more objects X/Y/Z that implement a ThirdTrait. You can make a struct that holds a generic of each type:

struct ComposeTraits<T1: SomeTrait, T2: OtherTrait, T3: ThirdTrait> {
    first: T1,
    second: T2,
    third: T3,
}

The goal being that if the ComposeTraits struct has functioning versions for each generic it can perform some desired outer/interfacing methods, an InterfaceTrait, while avoiding internal enums or Box<dyn InternalTrait>s such that everything inside the composed struct works together without indirection/matching.

With this basic setup, you could, if you wanted, match against a tuple of enums selecting for each type to be used, and in each match arm return a Box<dyn InterfaceTrait>. However, since this is composing multiple options, it then requires a number of unique match arms equal to the product of the number variants for each type (4x5x3 = 60 arms for the above setup, yikes!)

However, I've mocked up an idea that extends/reworks the builder pattern to be able to set each generic in turn (i.e. only matching against one enum per type to set), and still returns the final constructed object into a single, outer Box.

Here is it working for a simple A1/A2+B1/B2 setup, but it should be sufficient to show that it can be extended to any number of generics with any number of type options: playground

The major goals for this were:

  1. Create a single runtime-selected and fully-typed struct behind just one heap allocation.
  2. Avoiding using enums with extreme size differences for each type.
  3. Maintain internal impl Trait arguments so that all compiler optimizations can be applied.

I'm guessing this might be a bit niche, but I'd really appreciate any thoughts on this. Thanks!

2

If I wanted to implement this by your proposed method, I probably would just write out the 60 match arms, possibly via a custom macro_rules! helper. This will more clearly show what your code will compile to, as opposed to your method chaining approach where you see 2 or 3 methods but actually it compiles down to 60.

Theoretically you could imagine an implementation that doesn't involve all 60 arms in the final binary, instead storing three vtables and then the three structs in the allocation. Unfortunately, I don't think there is any way to implement that in safe Rust.

1 Like
3

Yeah, the direct match option (especially via macro) could be cleaner/clearer. But I was hoping to make this extensible so I could cap things off into a Box<dyn Trait> at various possible levels of assembly of these bubbles of interfacing impls so that I could try benchmarking and checking the size of the codebase. I also like the simplicity of the set_a/set_b methods for internal type overwrites. Though I suppose my aims might be possible to accomplish using those with macros... I'll try that next!

Can you give any hints about how you might go about this with unsafe? I'm certainly up for that if there are further simplifications that could be made.

4

If you insert a Box<dyn Trait> at each level, then you can avoid the 60 different functions in the executable, but you would no longer have all three structs in the same allocation.

Sure! There are a few variations on how to do it, but here is one possibility. Basically, the idea is to unsafely create the following struct at runtime:

// This struct doesn't exist in your code, but we pretend it does.
#[repr(C)]
struct AllocatedData<Ta, Tb, Tc> {
    header: Header,
    struct_a: Ta,
    struct_b: Tb,
    struct_c: Tc,
}

struct Header {
    dealloc_info: std::alloc::Layout,
    get_vtable_a: unsafe fn(*mut u8) -> *mut dyn TraitA,
    ptr_a: *mut u8,
    get_vtable_b: unsafe fn(*mut u8) -> *mut dyn TraitB,
    ptr_b: *mut u8,
    get_vtable_c: unsafe fn(*mut u8) -> *mut dyn TraitC,
    ptr_c: *mut u8,
}

Then, to mimic the AllocationData struct, we compute the struct's size and alignment, as well as the offset for each field in the struct using the #[repr(C)] algorithm:

fn fix_alignment(min_offset: usize, align: usize) -> usize {
    let mut offset = min_offset;
    while offset % align != 0 {
        offset += 1;
    }
    offset
}

let (size_a, align_a) = ...;
let (size_b, align_b) = ...;
let (size_c, align_c) = ...;

let offset_a = fix_alignment(std::mem::size_of::<Header>(), align_a);
let offset_b = fix_alignment(offset_a + size_a, align_b);
let offset_c = fix_alignment(offset_b + size_b, align_c);

let total_align = max(std::mem::align_of::<Header>(), align_a, align_b, align_c);
let total_size = offset_c + size_c;
let layout = std::alloc::Layout::from_size_align(total_size, total_align).unwrap();

let allocation = std::alloc::alloc(layout);

Now, we can create our three structs and write them into our new allocation:

// similarly for b and c
let ptr_a = allocation.add(offset_a);
match type_a {
    Type1 => std::ptr::write(ptr_a as *mut Type1, value_of_type1()),
    Type2 => std::ptr::write(ptr_a as *mut Type2, value_of_type2()),
    Type3 => std::ptr::write(ptr_a as *mut Type3, value_of_type3()),
}

Then we can write the header at offset zero:

let header = Header {
    dealloc_info: layout,
    get_vtable_a: ...,
    get_vtable_b: ...,
    get_vtable_c: ...,
    ptr_a: allocation.add(offset_a),
    ptr_b: allocation.add(offset_b),
    ptr_c: allocation.add(offset_c),
}
std::ptr::write(allocation as *mut Header, header);

Here, the way you get the vtable functions is this:

fn get_vtable_a<T: TraitA + 'static>(ptr: *mut u8) -> *mut dyn TraitA {
    let ptr = ptr as *mut T;
    ptr
}

let vtable_a = match type_a {
    Type1 => get_vtable_a::<Type1>(),
    Type2 => get_vtable_a::<Type2>(),
    Type3 => get_vtable_a::<Type3>(),
};

Now you have pretty much everything and you can put the pointer into a struct;

struct ComposeTraitBox {
   ptr: *mut Header, // or NonNull
}

Whenever you want to access one of the fields, you just pass the appropriate pointer in the header to the appropriate get_vtable function, and it returns a *mut dyn TraitA raw pointer, which you can turn into a reference.

To drop the box, you first drop the three values by getting the *mut dyn TraitA raw pointer and passing it to drop_in_place. Then after that, you use the dealloc_info in the Header struct together with the ComposeTraitBox::ptr field and pass them to the std::alloc::dealloc function.

3 Likes
5

Whoah! I think I'll let the compiler keep managing alignment and layout for now, but that's certainly an interesting path to consider. Thanks for such a detailed response! And hooray for new fun things to study!!

6

I came to a realization about how object safe traits can take self: Box<Self> and how to use that to manage a typed/typestated inner struct via trait access in a much simpler and more ergonomic way! I think this is nearly to a final usable form: playground link

7

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.