Inlined vtable pattern to replace `&mut FnMut(T) -> U` type

Sometimes I need to use type erasure to reduce binary size:

fn foo(f: &mut dyn FnMut(T) -> U) {
    ...
    // Code that calls `f`.
    ...
}

I am wondering if this can be further optimized. As far as I know, dyn Trait needs a vtable to work, and the vtable contains metadata like size, align, drop_in_place address, etc. But in this specific case, most of the metadata is not used, so I imagine a wrapper type that behaves like a &mut dyn FnMut(Foo) -> Bar, but without any unnecessary vtable stuff, something like this:

pub struct RefFnMut1<'a, T, U> {
    // The address the `FnMut` object,
    data: NonNull<()>,
    // The address of a wrapper function that calls `FnMut::call_mut` function.
    call_mut_fn: unsafe fn(NonNull<()>, T) -> U,
    _phantom: PhantomData<&'a mut ()>,
}

impl<'a, T, U> RefFnMut1<'a, T, U> {
    pub fn new<F>(f: &'a mut F) -> Self
    where
        F: FnMut(T) -> U,
    {
        ...
    }

    pub fn call_mut(&mut self, arg: T) -> U {
        unsafe { (self.call_mut_fn)(self.data, arg) }
    }
}

Here is a proof of concept playground link: https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=ba86efedf940958705c3acc06c908fff.

Basically, this is like a inlined vtable, but only contains useful information.

Is there any problem with this pattern? Is there any existing crate that implements this pattern? I am not very confident with writing unsafe codes, so an existing implementation will be extremely helpful.

For what it's worth, I don't see any issues in your implementation. The only unsafe part is erasing the type and lifetime from data, and you correctly handle it through your 'a lifetime and call_mut_fn::<F, T> pointer. There also don't appear to be any variance issues, since RefMut1 is implicitly invariant over F, and correctly contravariant over T and covariant over U.

1 Like

Note that &'a mut FnMut(T) -> U is actually invariant over T and U though. I would need to think more about whether relaxing this can pose a problem here.

It shouldn't be a problem, since the call_mut_fn() wrapper naturally performs the proper upcasting for the argument and return type. Indeed, a trivial variation of this without the type erasure for F can be written entirely in safe code (Rust Playground). This variation is explicitly invariant over F.

1 Like

For this to work the compile must make sure that neither of size_of_val(), align_of_val() or drop_in_place() are called for the value, including nested calls. I don't think we have the machinery for that.

If anyone is interested, I just published a crate based on this idea: https://docs.rs/simple-ref-fn/latest/simple_ref_fn/.