ffi_trait 1.2.0

#![doc = include_str!("../README.md")]
#![no_std]

use alloc::{alloc::dealloc, boxed::Box};
use core::{alloc::Layout, marker::PhantomData, mem::ManuallyDrop, ptr::NonNull};

extern crate alloc;

/// # Safety
/// `Self` must be layout compatible with [`DropVtable`]
pub unsafe trait VtableFor<T>: Copy + 'static {
    const VTABLE: &'static Self;
}

#[repr(C)]
pub struct RefDyn<'a, Vtable: 'static> {
    ptr: NonNull<()>,
    vtable: &'static Vtable,
    _ref: PhantomData<&'a ()>,
}

impl<'a, Vtable: 'static> RefDyn<'a, Vtable> {
    pub const fn new<T>(value: &'a T) -> Self
    where
        Vtable: VtableFor<T>,
    {
        Self {
            ptr: NonNull::new(value as *const _ as _).unwrap(),
            vtable: Vtable::VTABLE,
            _ref: PhantomData,
        }
    }

    /// # Safety
    /// `ptr` must be valid to use with `vtable`
    pub const unsafe fn from_ptr_vtable_unchecked(
        ptr: NonNull<()>,
        vtable: &'static Vtable,
    ) -> Self {
        Self {
            ptr,
            vtable,
            _ref: PhantomData,
        }
    }

    /// # Safety
    /// `ptr` must point to a valid `T`
    pub const unsafe fn get_unchecked<T>(&self) -> &T {
        unsafe { self.ptr.cast().as_ref() }
    }

    /// # Safety
    /// `ptr` must point to a valid `T`
    pub const unsafe fn get_mut_unchecked<T>(&mut self) -> &mut T {
        unsafe { self.ptr.cast().as_mut() }
    }

    pub const fn vtable(&self) -> &'static Vtable {
        self.vtable
    }

    pub const fn as_ptr(&self) -> NonNull<()> {
        self.ptr
    }
}

#[repr(C)]
pub struct OwnDyn<Vtable: 'static> {
    ptr: NonNull<()>,
    vtable: &'static Vtable,
}

impl<Vtable: 'static> OwnDyn<Vtable> {
    pub fn new<T>(value: Box<T>) -> Self
    where
        Vtable: VtableFor<T>,
    {
        Self {
            ptr: NonNull::new(Box::into_raw(value).cast()).unwrap(),
            vtable: Vtable::VTABLE,
        }
    }

    /// # Safety
    /// `ptr` must be allocated with the global allocator and must be valid to use with `vtable`
    pub const unsafe fn from_ptr_vtable_unchecked(
        ptr: NonNull<()>,
        vtable: &'static Vtable,
    ) -> Self {
        Self { ptr, vtable }
    }

    /// # Safety
    /// `ptr` must point to a valid `T`
    pub unsafe fn into_box<T>(self) -> Box<T> {
        let this = ManuallyDrop::new(self);
        unsafe { Box::from_raw(this.ptr.as_ptr().cast()) }
    }

    /// # Safety
    /// `ptr` must point to a valid `T`
    pub const unsafe fn get_unchecked<T>(&self) -> &T {
        unsafe { self.ptr.cast().as_ref() }
    }

    /// # Safety
    /// `ptr` must point to a valid `T`
    pub const unsafe fn get_mut_unchecked<T>(&mut self) -> &mut T {
        unsafe { self.ptr.cast().as_mut() }
    }

    pub const fn vtable(&self) -> &'static Vtable {
        self.vtable
    }

    pub const fn as_ptr(&self) -> NonNull<()> {
        self.ptr
    }
}

impl<Vtable: 'static> Drop for OwnDyn<Vtable> {
    fn drop(&mut self) {
        unsafe {
            let drop_vtable = &*(&raw const self.vtable as *const DropVtable);
            (drop_vtable.drop)(self.ptr);
            dealloc(
                self.ptr.as_ptr().cast(),
                Layout::from_size_align_unchecked(drop_vtable.size, drop_vtable.align),
            );
        }
    }
}

#[derive(Clone, Copy)]
#[repr(C)]
pub struct DropVtable {
    pub size: usize,
    pub align: usize,
    pub drop: unsafe extern "C" fn(NonNull<()>),
}

unsafe impl<T> VtableFor<T> for DropVtable {
    const VTABLE: &'static Self = {
        unsafe extern "C" fn drop_thunk<T>(ptr: NonNull<()>) {
            unsafe { ptr.cast::<T>().drop_in_place() }
        }

        &Self {
            size: size_of::<T>(),
            align: align_of::<T>(),
            drop: drop_thunk::<T>,
        }
    };
}

#[doc(hidden)]
pub mod __macro_helper {
    pub use core::ptr::NonNull;
}

#[macro_export]
macro_rules! vtable_for_trait {
    {
        $trait_vis:vis trait $trait_name:ident {
            $(fn $method_name:ident(& $(mut $(@$self_mut:tt)?)? self $(, $param_name:ident : $param_type:ty)* $(,)?) $(-> $return_type:ty)?;)*
        }

        $vtable_vis:vis vtable $vtable_name:ident;
    } => {
        $trait_vis trait $trait_name {
            $(fn $method_name(& $(mut $(@$self_mut)?)? self $(, $param_name : $param_type)*) $(-> $return_type)?;)*
        }

        #[derive(Clone, Copy)]
        #[repr(C)]
        $vtable_vis struct $vtable_name {
            pub drop_vtable: $crate::DropVtable,
            $(pub $method_name: unsafe extern "C" fn($crate::__macro_helper::NonNull<()> $(, $param_type)*) $(-> $return_type)?,)*
        }

        unsafe impl<__T: $trait_name> $crate::VtableFor<__T> for $vtable_name {
            const VTABLE: &'static Self = {
                $(
                    unsafe extern "C" fn $method_name<__T: $trait_name>(__ptr: $crate::__macro_helper::NonNull<()> $(, $param_name : $param_type)*) $(-> $return_type)? {
                        let __ref = unsafe { & $(mut $(@$self_mut)?)? *__ptr.cast::<__T>().as_ptr() };
                        <__T as $trait_name>::$method_name(__ref $(,$param_name)*)
                    }
                )*

                &Self {
                    drop_vtable: *<$crate::DropVtable as $crate::VtableFor<__T>>::VTABLE,
                    $(
                        $method_name: $method_name::<__T>,
                    )*
                }
            };
        }

        impl<'a> $trait_name for $crate::RefDyn<'a, $vtable_name> {
            $(
                fn $method_name(& $(mut $(@$self_mut)?)? self $(, $param_name : $param_type)*) $(-> $return_type)? {
                    unsafe { (self.vtable().$method_name)(self.as_ptr() $(, $param_name)*) }
                }
            )*
        }

        impl<'a> $trait_name for $crate::OwnDyn<$vtable_name> {
            $(
                fn $method_name(& $(mut $(@$self_mut)?)? self $(, $param_name : $param_type)*) $(-> $return_type)? {
                    unsafe { (self.vtable().$method_name)(self.as_ptr() $(, $param_name)*) }
                }
            )*
        }
    };
}

#[cfg(test)]
mod tests {
    use crate::RefDyn;

    vtable_for_trait! {
        trait Bar {
            fn foo(&self, x: i32) -> f32;
            fn bar(&mut self);
        }

        vtable BarVtable;
    }

    #[test]
    fn bar() {
        struct Test;

        impl Bar for Test {
            fn foo(&self, x: i32) -> f32 {
                x as f32
            }

            fn bar(&mut self) {}
        }

        let t = Test;
        let bar = RefDyn::<BarVtable>::new(&t);
        assert_eq!(bar.foo(5), 5.0);
    }
}