wasmtime-module-dsl.rs

use cranelift_codegen::ir::types;
use cranelift_codegen::{ir, isa};
use cranelift_entity::PrimaryMap;
use cranelift_wasm::{DefinedFuncIndex, Global as WasmGlobal, GlobalInit};
use hlist::{Cons, Here, Nil, There};
use std::any::Any;
use std::cell::RefCell;
use std::collections::HashMap;
use std::marker::PhantomData;
use std::rc::Rc;
use target_lexicon::HOST;
use wasmtime_environ::{translate_signature, Export, MemoryPlan, Module, TablePlan};
use wasmtime_jit::target_tunables;
use wasmtime_runtime::{Imports, InstanceHandle, InstantiationError, VMContext, VMFunctionBody};

pub use cranelift_wasm::{Memory, Table, TableElementType};
pub use hlist;

pub struct Builder<'a> {
    pub module: Module,
    pub functions: PrimaryMap<DefinedFuncIndex, *const VMFunctionBody>,
    pub triple: &'a target_lexicon::Triple,
}

impl<'a> Builder<'a> {
    pub fn new(triple: &'a target_lexicon::Triple) -> Self {
        let module = Module::new();
        let functions: PrimaryMap<DefinedFuncIndex, *const VMFunctionBody> = PrimaryMap::new();

        Builder {
            module,
            functions,
            triple,
        }
    }

    pub fn instantiate_module<T: Any>(
        self,
        host_data: T,
    ) -> Result<InstanceHandle, InstantiationError> {
        let imports = Imports::none();
        let data_initializers = Vec::new();
        let signatures = PrimaryMap::new();

        InstanceHandle::new(
            Rc::new(self.module),
            Rc::new(RefCell::new(HashMap::new())),
            self.functions.into_boxed_slice(),
            imports,
            &data_initializers,
            signatures.into_boxed_slice(),
            None,
            Box::new(host_data),
        )
    }
}

pub enum Value {
    F64(u64),
    F32(u32),
    I32(u32),
    I64(u64),
}

impl From<f64> for Value {
    fn from(other: f64) -> Self {
        Value::F64(other.to_bits())
    }
}
impl From<f32> for Value {
    fn from(other: f32) -> Self {
        Value::F32(other.to_bits())
    }
}
impl From<i32> for Value {
    fn from(other: i32) -> Self {
        Value::I32(other as _)
    }
}
impl From<i64> for Value {
    fn from(other: i64) -> Self {
        Value::I64(other as _)
    }
}
impl From<u32> for Value {
    fn from(other: u32) -> Self {
        Value::I32(other)
    }
}
impl From<u64> for Value {
    fn from(other: u64) -> Self {
        Value::I64(other)
    }
}

pub trait Exports<Counter, HostData>: HasHostData {
    fn build(self, builder: &mut Builder<'_>) -> Self::HostData;
}

pub struct ExportDef<T> {
    name: String,
    val: T,
}

#[derive(Debug, PartialEq, Eq)]
pub enum Mutability {
    Mutable,
    Immutable,
}

impl Default for Mutability {
    fn default() -> Self {
        Mutability::Immutable
    }
}

pub struct Func<T>(T);
pub struct Global<T>(T, Mutability);

pub struct BindArgs<F, A>(F, PhantomData<A>);

pub trait BindArgType: Sized {
    fn bind<A>(self) -> BindArgs<Self, A>;
}

impl<T> BindArgType for T {
    fn bind<A>(self) -> BindArgs<Self, A> {
        BindArgs(self, PhantomData)
    }
}

pub trait IntoSigTypes {
    fn sig_types() -> Vec<ir::AbiParam>;
}

pub trait IntoSigType {
    fn sig_type() -> types::Type;
}

impl IntoSigType for i32 {
    fn sig_type() -> types::Type {
        types::I32
    }
}

impl IntoSigType for u32 {
    fn sig_type() -> types::Type {
        types::I32
    }
}

impl IntoSigType for i64 {
    fn sig_type() -> types::Type {
        types::I64
    }
}

impl IntoSigType for u64 {
    fn sig_type() -> types::Type {
        types::I64
    }
}

impl IntoSigType for f32 {
    fn sig_type() -> types::Type {
        types::F32
    }
}

impl IntoSigType for f64 {
    fn sig_type() -> types::Type {
        types::F64
    }
}

pub trait MkTrampoline<A, O> {
    fn make_trampoline<F, C, HostData>() -> *const VMFunctionBody
    where
        F: CallMut<Args = A, Output = O>,
        HostData: hlist::Find<Func<F>, C> + 'static;
}

/// Version of `FnMut` that constrains the arguments to only a single type
pub trait CallMut {
    type Args;
    type Output;

    fn call_mut(&mut self, args: Self::Args) -> Self::Output;
}

macro_rules! impl_functionlike_traits {
    ($first_a:ident $(, $rest_a:ident)*) => {
        impl<$first_a $(, $rest_a)*, __O> MkTrampoline<($first_a, $($rest_a),*), __O> for
            (($first_a, $($rest_a),*), __O)
        {
            fn make_trampoline<__F, __C, HostData>() -> *const VMFunctionBody
            where
                __F: CallMut<Args = ($first_a, $($rest_a),*), Output = __O>,
                HostData: hlist::Find<Func<__F>, __C> + 'static,
            {
                (|vmctx: &mut VMContext, $first_a: $first_a $(, $rest_a: $rest_a)*| {
                    let func = &mut hlist::Find::<Func<__F>, __C>::get_mut(
                        unsafe { vmctx.host_state() }
                            .downcast_mut::<HostData>()
                            .expect("Programmer error: Invalid host data"),
                    )
                    .0;
                    func.call_mut(($first_a, $($rest_a),*))
                }) as fn(&mut VMContext, $first_a, $($rest_a),*) -> __F::Output as _
            }
        }

        impl<__F, $first_a $(, $rest_a)*, __O> CallMut for BindArgs<__F, ($first_a, $($rest_a),*)>
        where
            __F: FnMut($first_a, $($rest_a),*) -> __O,
        {
            type Args = ($first_a, $($rest_a),*);
            type Output = __O;

            fn call_mut(&mut self, args: Self::Args) -> Self::Output {
                let ($first_a, $($rest_a),*) = args;
                (self.0)($first_a $(, $rest_a)*)
            }
        }

        impl<$first_a $(, $rest_a)*> IntoSigTypes for ($first_a, $($rest_a),*)
        where
            $first_a: IntoSigType
            $(, $rest_a: IntoSigType)*
        {
            fn sig_types() -> Vec<ir::AbiParam> {
                vec![ir::AbiParam::new($first_a::sig_type())$(, ir::AbiParam::new($rest_a::sig_type()))*]
            }
        }

        impl_functionlike_traits!($($rest_a),*);
    };
    () => {
        impl<__O> MkTrampoline<(), __O> for ((), __O) {
            fn make_trampoline<F, C, HostData>() -> *const VMFunctionBody
            where
                F: CallMut<Args = (), Output = __O>,
                HostData: hlist::Find<Func<F>, C> + 'static,
            {
                (|vmctx: &mut VMContext| {
                    let func = &mut hlist::Find::<Func<F>, C>::get_mut(
                        unsafe { vmctx.host_state() }
                            .downcast_mut::<HostData>()
                            .expect("Programmer error: Invalid host data"),
                    )
                    .0;
                    func.call_mut(())
                }) as fn(&mut VMContext) -> F::Output as _
            }
        }

        impl<F, __O> CallMut for BindArgs<F, ()>
        where
            F: FnMut() -> __O,
        {
            type Args = ();
            type Output = __O;

            fn call_mut(&mut self, _args: ()) -> Self::Output {
                (self.0)()
            }
        }

        impl IntoSigTypes for () {
            fn sig_types() -> Vec<ir::AbiParam> {
                vec![]
            }
        }
    };
}

impl<T> IntoSigTypes for T
where
    T: IntoSigType,
{
    fn sig_types() -> Vec<ir::AbiParam> {
        vec![ir::AbiParam::new(T::sig_type())]
    }
}

#[allow(non_snake_case)]
mod dummy {
    use super::*;

    impl_functionlike_traits!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P);
}

pub trait HasHostData {
    type HostData;
}

impl<F, Rest> HasHostData for Cons<ExportDef<Func<F>>, Rest>
where
    Rest: HasHostData,
{
    type HostData = Cons<Func<F>, <Rest as HasHostData>::HostData>;
}

impl<T, Rest> HasHostData for Cons<ExportDef<Global<T>>, Rest>
where
    Rest: HasHostData,
{
    type HostData = <Rest as HasHostData>::HostData;
}

impl<Rest> HasHostData for Cons<ExportDef<Memory>, Rest>
where
    Rest: HasHostData,
{
    type HostData = <Rest as HasHostData>::HostData;
}

impl<Rest> HasHostData for Cons<ExportDef<Table>, Rest>
where
    Rest: HasHostData,
{
    type HostData = <Rest as HasHostData>::HostData;
}

impl HasHostData for Nil {
    type HostData = Nil;
}

impl<F, Rest, Counter, HostData> Exports<Counter, HostData> for Cons<ExportDef<Func<F>>, Rest>
where
    Cons<ExportDef<Func<F>>, Rest>:
        HasHostData<HostData = Cons<Func<F>, <Rest as HasHostData>::HostData>>,
    F: CallMut,
    F::Args: IntoSigTypes,
    F::Output: IntoSigTypes,
    (F::Args, F::Output): MkTrampoline<F::Args, F::Output>,
    Rest: Exports<There<Counter>, HostData>,
    HostData: hlist::Find<Func<F>, Counter> + 'static,
{
    fn build(self, builder: &mut Builder<'_>) -> Self::HostData {
        let calling_convention = isa::CallConv::triple_default(&builder.triple);
        let pointer_type = types::Type::triple_pointer_type(&builder.triple);
        let sig = builder.module.signatures.push(translate_signature(
            ir::Signature {
                params: F::Args::sig_types(),
                returns: F::Output::sig_types(),
                call_conv: calling_convention,
            },
            pointer_type,
        ));

        let func = builder.module.functions.push(sig);
        builder
            .module
            .exports
            .insert(self.0.name, Export::Function(func));
        builder
            .functions
            .push(<(F::Args, F::Output)>::make_trampoline::<
                F,
                Counter,
                HostData,
            >());

        Cons(self.0.val, self.1.build(builder))
    }
}

impl<T, Rest, Counter, HostData> Exports<Counter, HostData> for Cons<ExportDef<Global<T>>, Rest>
where
    Cons<ExportDef<Global<T>>, Rest>: HasHostData<HostData = <Rest as HasHostData>::HostData>,
    Rest: Exports<Counter, HostData>,
    T: Into<Value>,
{
    fn build(self, builder: &mut Builder<'_>) -> Self::HostData {
        let Global(value, mutability) = self.0.val;
        let global = match value.into() {
            Value::F32(val) => builder.module.globals.push(WasmGlobal {
                ty: types::F32,
                mutability: mutability == Mutability::Mutable,
                initializer: GlobalInit::F32Const(val as _),
            }),
            Value::F64(val) => builder.module.globals.push(WasmGlobal {
                ty: types::F64,
                mutability: mutability == Mutability::Mutable,
                initializer: GlobalInit::F64Const(val as _),
            }),
            Value::I32(val) => builder.module.globals.push(WasmGlobal {
                ty: types::I32,
                mutability: mutability == Mutability::Mutable,
                initializer: GlobalInit::I32Const(val as _),
            }),
            Value::I64(val) => builder.module.globals.push(WasmGlobal {
                ty: types::I64,
                mutability: mutability == Mutability::Mutable,
                initializer: GlobalInit::I64Const(val as _),
            }),
        };
        builder
            .module
            .exports
            .insert(self.0.name, Export::Global(global));

        self.1.build(builder)
    }
}

impl<Rest, Counter, HostData> Exports<Counter, HostData> for Cons<ExportDef<Memory>, Rest>
where
    Cons<ExportDef<Memory>, Rest>: HasHostData<HostData = <Rest as HasHostData>::HostData>,
    Rest: Exports<Counter, HostData>,
{
    fn build(self, builder: &mut Builder<'_>) -> Self::HostData {
        let tunables = target_tunables(builder.triple);
        let memory = builder
            .module
            .memory_plans
            .push(MemoryPlan::for_memory(self.0.val, &tunables));
        builder
            .module
            .exports
            .insert(self.0.name, Export::Memory(memory));

        self.1.build(builder)
    }
}

impl<Rest, Counter, HostData> Exports<Counter, HostData> for Cons<ExportDef<Table>, Rest>
where
    Cons<ExportDef<Table>, Rest>: HasHostData<HostData = <Rest as HasHostData>::HostData>,
    Rest: Exports<Counter, HostData>,
{
    fn build(self, builder: &mut Builder<'_>) -> Self::HostData {
        let tunables = target_tunables(builder.triple);
        let memory = builder
            .module
            .table_plans
            .push(TablePlan::for_table(self.0.val, &tunables));
        builder
            .module
            .exports
            .insert(self.0.name, Export::Table(memory));

        self.1.build(builder)
    }
}

impl<Counter, HostData> Exports<Counter, HostData> for Nil {
    fn build(self, _builder: &mut Builder<'_>) -> Self::HostData {
        Nil
    }
}

pub trait Instantiate {
    fn instantiate(self) -> Result<InstanceHandle, InstantiationError>;
}

impl<T> Instantiate for T
where
    T: HasHostData,
    <T as HasHostData>::HostData: 'static,
    T: Exports<Here, <T as HasHostData>::HostData>,
{
    fn instantiate(self) -> Result<InstanceHandle, InstantiationError> {
        let mut builder = Builder::new(&HOST);
        let data = self.build(&mut builder);

        builder.instantiate_module(data)
    }
}

#[macro_export]
macro_rules! exports {
    ($name:ident: $val:expr $(, $k:ident: $v:expr)* $(,)*) => {{
        $crate::hlist::Cons($crate::ExportDef {
            name: stringify!($name).to_owned(),
            val: $val,
        }, exports!($($k:$v),*))
    }};
    () => {
        $crate::hlist::Nil
    }
}

#[cfg(test)]
mod tests {
    use super::{BindArgType, Func, Global, Instantiate, Memory, Table, TableElementType};

    #[test]
    fn example() {
        fn hello_world() {
            println!("Hello, world!");
        }

        fn print_and_return(val: i32) -> i32 {
            println!("{}", val);
            val
        }

        let mut counter = 100;
        let my_closure = move |inc: u32| -> u32 {
            counter += inc;
            counter
        };

        (exports! {
            do_thing: Func(hello_world.bind::<()>()),
            my_glob: Global(100u64, Default::default()),
            print_and_return: Func(print_and_return.bind::<(i32,)>()),
            counting_func: Func(my_closure.bind::<(u32,)>()),
            memory: Memory {
                minimum: 1,
                maximum: Some(2),
                shared: false,
            },
            table: Table {
                ty: TableElementType::Func,
                minimum: 10,
                maximum: Some(20),
            },
        })
        .instantiate()
        .unwrap();
    }
}