melior icon indicating copy to clipboard operation
melior copied to clipboard

[Feature request] Wrapper type for callable JIT functions.

Open BadBastion opened this issue 11 months ago • 0 comments

Wanted to get thoughts on adding a JitFn wrapper to make ExecutionEngine calls easier.

Pseudo(-ish) code example:

// Based on inkwells JitFn
// https://github.com/TheDan64/inkwell/blob/master/src/execution_engine.rs#L508
pub struct JitFn<'a, F> {
    engine: &'a ExecutionEngine,
    inner: F,
}

impl<'a, F> JitFn<'a, F>
where
    F: UnsafeFunctionPointer,
{
    fn new(engine: &'a ExecutionEngine, func_name: &str) -> Result<Self, LookupError> {
        let address = engine.lookup(func_name);
        if address.is_null() {
            Err(LookupError::new(func_name))
        } else {
            let inner = unsafe { transmute_copy(&address) };
            Ok(Self { engine, inner })
        }
    }
}

/// Marker trait representing an unsafe function pointer (`unsafe extern "C" fn(A, B, ...) -> Output`).
pub trait UnsafeFunctionPointer: private::SealedUnsafeFunctionPointer {}
impl<F: private::SealedUnsafeFunctionPointer> UnsafeFunctionPointer for F {}

mod private {
    /// A sealed trait which ensures nobody outside this crate can implement
    /// `UnsafeFunctionPointer`.
    ///
    /// See https://rust-lang-nursery.github.io/api-guidelines/future-proofing.html
    pub trait SealedUnsafeFunctionPointer: Copy {}
}

macro_rules! impl_unsafe_fn {
    (@recurse $first:ident $( , $rest:ident )*) => {
        impl_unsafe_fn!($( $rest ),*);
    };

    (@recurse) => {};

    ($( $param:ident ),*) => {
        impl<Output, $( $param ),*> private::SealedUnsafeFunctionPointer for unsafe extern "C" fn($( $param ),*) -> Output {}

        impl<'a, Output, $( $param ),*> JitFunction<'a, unsafe extern "C" fn($( $param ),*) -> Output> {
            #[allow(non_snake_case)]
            #[inline(always)]
            pub unsafe fn call(&self, $( mut $param: $param ),*) -> Output {
                (self.inner)($( $param ),*)
            }
        }

        impl_unsafe_fn!(@recurse $( $param ),*);
    };
}
// Recursively implement the trait for each parameter count
impl_unsafe_fn!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P);

BadBastion avatar Mar 14 '24 17:03 BadBastion