🤓.rs

//! Warning! This code is probably overcomplicated. I'm questioning if specifically the size
//! tests are worth the complexity. I decided "yes" because its generic enough that it can
//! be turned into a library and used throughout the codebase anywhere where ffi-safe closures
//! must be used.
//!
//! There are two main sections:
//! - The [`make_ffi_fn`] and [`ffi_fn`] functions
//! - Compile-time checks to make sure that trait object vtables are not being sliced off when we
//!   cast to a pointer.
//!
//! The latter employs some very cursed typesystem hacks to get around several Rust limitations. In particular,
//! to get around the inability for a const to access the generics of its enclosing function.
//!
//! For explanation about why trait objects are not ffi-safe and how to solve it, read this:
//! https://adventures.michaelfbryan.com/posts/rust-closures-in-ffi/
//! The TLDR is that trait objects are not FFI safe because they are fat pointers. So we can't use a Box<dyn
//! Trait> in the FFI boundary. This implies whatever function we give C, must already know the concrete
//! closure type. We can accomplish this by "instantiating" the generic `ffi_fn<F: FnMut>` function
//! with the concrete closure type, and then using this new monomorphized function in FFI.
//!
//! You can also look at this playground link I made for some proof that trait objects and regular
//! references have different sizes and can't just be turned into c style pointers trivially:
//! https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=896ace6ed2e1abd64397a210c375f81a
//!
//! The typesystem abuse comes in part from here, but I made heavy modifications because the original
//! approach didn't actually work:
//! https://users.rust-lang.org/t/is-it-possible-to-assert-at-compile-time-that-foo-t-is-not-called-with-a-zst/67685/3

extern crate alloc;
use alloc::boxed::Box;
use core::ffi::c_void;

/// Monomorphizes `ffi_fn` returning a new `fn` that we can use in FFI as well as the `*mut c_void`
/// that should be passed to it.
fn make_ffi_fn<Args, F: FnMut(Args) + 'static + ?Sized>(
    closure: Box<F>,
) -> (unsafe extern "C" fn(Args, *mut c_void), *mut c_void) {
    use self::cursed_typesystem_abuse::PanicWhenFat;

    let closure: Box<F> = PanicWhenFat::check(closure);
    let ptr: *mut F = Box::into_raw(closure);
    let ptr = ptr as *mut c_void;

    // monomorphizes your function 😳
    (ffi_fn::<Args, Box<F>>, ptr)
}

// This function is generic, to use it in FFI we must monomorphize it with [`make_ffi_fn`].
unsafe extern "C" fn ffi_fn<Args, F: FnMut(Args)>(args: Args, closure: *mut c_void) {
    let closure: &mut F = unsafe {
        // Right here, we cast from the void back to a known concrete closure type!
        (closure as *mut F)
            .as_mut()
            .expect("Got a null pointer, this should not have been possible.")
    };

    closure(args)
}

/// This module is used purely for checking at compile time if the the closure is a fat pointer or
/// not.
mod cursed_typesystem_abuse {
    const fn assert_same_size_as_pointer<T>() {
        if core::mem::size_of::<T>() == core::mem::size_of::<usize>() {
            return;
        }
        panic!("`T` was not the same size as a pointer!!!");
    }

    /// Abuses typesystem to run a const fn that gets passed a generic type. If you just try to
    /// make a regular function for this, you'll get an error about not being able to use generics
    /// from the enclosing function.
    pub trait PanicWhenFat: Sized {
        const CHECK: () = assert_same_size_as_pointer::<Self>();

        /// Checks that Self has the exact same size as a pointer. Useful for asserting against fat
        /// pointers. Otherwise, it is just the identity function.
        fn check(self) -> Self {
            let _ = Self::CHECK;
            self
        }

        /// Same as `check` but doesn't consume the value.
        fn check_ref(&self) {
            let _ = Self::CHECK;
        }
    }
    impl<T> PanicWhenFat for T {}
}

// Without a public function using `make_ffi_fn`, the compiler may never choose to monomorphize and
// we won't see the result of `PanicWhenFat::check`. By providing these tests, we ensure there are
// concrete monomorphizations of the function that will definitely be run by the compiler.
#[cfg(test)]
mod tests {
    use super::*;

    extern crate alloc;
    use alloc::boxed::Box;

    // If you remove `#[test]` the compiler will strip out this function and the checks will never
    // run.
    #[test]
    pub fn test_ffi_fn_size_checks() {
        use cursed_typesystem_abuse::PanicWhenFat;

        let monomorphized = Box::new(|_: ()| {});
        let trait_obj: Box<dyn FnMut(())> = Box::new(|_| {});

        PanicWhenFat::check_ref(&monomorphized);
        // PanicWhenFat::check_ref(&trait_obj); // This will fail

        // This will fail, because trait objects are larger than a pointer.
        // let (_, _) = make_ffi_fn(trait_obj);
        // This will work, because a box of a trait object is the size of a regular box.
        let (_, _) = make_ffi_fn(Box::new(trait_obj));
        // This will work, because its a regular box.
        let (_, _) = make_ffi_fn(monomorphized);
    }
}