use proc_macro::TokenStream;
use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
use quote::{format_ident, quote};
use syn::parse_macro_input;
use wiggle_generate::Names;
mod config;
use config::{AsyncConf, Asyncness, ModuleConf, TargetConf};
#[proc_macro]
pub fn wasmtime_integration(args: TokenStream) -> TokenStream {
let config = parse_macro_input!(args as config::Config);
let doc = config.load_document();
let names = Names::new(quote!(wasmtime_wiggle));
let modules = config.modules.iter().map(|(name, module_conf)| {
let module = doc
.module(&witx::Id::new(name))
.unwrap_or_else(|| panic!("witx document did not contain module named '{}'", name));
generate_module(
&module,
&module_conf,
&names,
&config.target,
&config.ctx.name,
&config.async_,
)
});
quote!( #(#modules)* ).into()
}
fn generate_module(
module: &witx::Module,
module_conf: &ModuleConf,
names: &Names,
target_conf: &TargetConf,
ctx_type: &syn::Type,
async_conf: &AsyncConf,
) -> TokenStream2 {
let fields = module.funcs().map(|f| {
let name_ident = names.func(&f.name);
quote! { pub #name_ident: wasmtime::Func }
});
let get_exports = module.funcs().map(|f| {
let func_name = f.name.as_str();
let name_ident = names.func(&f.name);
quote! { #func_name => Some(&self.#name_ident) }
});
let ctor_fields = module.funcs().map(|f| names.func(&f.name));
let module_name = module.name.as_str();
let linker_add = module.funcs().map(|f| {
let func_name = f.name.as_str();
let name_ident = names.func(&f.name);
quote! {
linker.define(#module_name, #func_name, self.#name_ident.clone())?;
}
});
let target_path = &target_conf.path;
let module_id = names.module(&module.name);
let target_module = quote! { #target_path::#module_id };
let mut fns = Vec::new();
let mut ctor_externs = Vec::new();
let mut host_funcs = Vec::new();
for f in module.funcs() {
let asyncness = async_conf.is_async(module.name.as_str(), f.name.as_str());
match asyncness {
Asyncness::Blocking => {}
Asyncness::Async => {
assert!(
cfg!(feature = "async"),
"generating async wasmtime Funcs requires cargo feature \"async\""
);
}
_ => {}
}
generate_func(
&module_id,
&f,
names,
&target_module,
ctx_type,
asyncness,
&mut fns,
&mut ctor_externs,
&mut host_funcs,
);
}
let type_name = module_conf.name.clone();
let type_docs = module_conf
.docs
.as_ref()
.map(|docs| quote!( #[doc = #docs] ))
.unwrap_or_default();
let constructor_docs = format!(
"Creates a new [`{}`] instance.
External values are allocated into the `store` provided and
configuration of the instance itself should be all
contained in the `cx` parameter.",
module_conf.name.to_string()
);
let config_adder_definitions = host_funcs.iter().map(|(func_name, body)| {
let adder_func = format_ident!("add_{}_to_config", names.func(&func_name));
let docs = format!(
"Add the host function for `{}` to a config under a given module and field name.",
func_name.as_str()
);
quote! {
#[doc = #docs]
pub fn #adder_func(config: &mut wasmtime::Config, module: &str, field: &str) {
#body
}
}
});
let config_adder_invocations = host_funcs.iter().map(|(func_name, _body)| {
let adder_func = format_ident!("add_{}_to_config", names.func(&func_name));
let module = module.name.as_str();
let field = func_name.as_str();
quote! {
Self::#adder_func(config, #module, #field);
}
});
quote! {
#type_docs
pub struct #type_name {
#(#fields,)*
}
impl #type_name {
#[doc = #constructor_docs]
pub fn new(store: &wasmtime::Store, ctx: std::rc::Rc<std::cell::RefCell<#ctx_type>>) -> Self {
#(#ctor_externs)*
Self {
#(#ctor_fields,)*
}
}
pub fn get_export(&self, name: &str) -> Option<&wasmtime::Func> {
match name {
#(#get_exports,)*
_ => None,
}
}
pub fn add_to_linker(&self, linker: &mut wasmtime::Linker) -> anyhow::Result<()> {
#(#linker_add)*
Ok(())
}
pub fn add_to_config(config: &mut wasmtime::Config) {
#(#config_adder_invocations)*
}
#(#config_adder_definitions)*
pub fn set_context(store: &wasmtime::Store, ctx: #ctx_type) -> Result<(), #ctx_type> {
store.set(std::rc::Rc::new(std::cell::RefCell::new(ctx))).map_err(|ctx| {
match std::rc::Rc::try_unwrap(ctx) {
Ok(ctx) => ctx.into_inner(),
Err(_) => unreachable!(),
}
})
}
#(#fns)*
}
}
}
fn generate_func(
module_ident: &Ident,
func: &witx::InterfaceFunc,
names: &Names,
target_module: &TokenStream2,
ctx_type: &syn::Type,
asyncness: Asyncness,
fns: &mut Vec<TokenStream2>,
ctors: &mut Vec<TokenStream2>,
host_funcs: &mut Vec<(witx::Id, TokenStream2)>,
) {
let rt = names.runtime_mod();
let name_ident = names.func(&func.name);
let (params, results) = func.wasm_signature();
let arg_names = (0..params.len())
.map(|i| Ident::new(&format!("arg{}", i), Span::call_site()))
.collect::<Vec<_>>();
let arg_decls = params
.iter()
.enumerate()
.map(|(i, ty)| {
let name = &arg_names[i];
let wasm = names.wasm_type(*ty);
quote! { #name: #wasm }
})
.collect::<Vec<_>>();
let ret_ty = match results.len() {
0 => quote!(()),
1 => names.wasm_type(results[0]),
_ => unimplemented!(),
};
let async_ = if asyncness.is_sync() {
quote!()
} else {
quote!(async)
};
let await_ = if asyncness.is_sync() {
quote!()
} else {
quote!(.await)
};
let runtime = names.runtime_mod();
let fn_ident = format_ident!("{}_{}", module_ident, name_ident);
fns.push(quote! {
#async_ fn #fn_ident(caller: &wasmtime::Caller<'_>, ctx: &mut #ctx_type #(, #arg_decls)*) -> Result<#ret_ty, wasmtime::Trap> {
unsafe {
let mem = match caller.get_export("memory") {
Some(wasmtime::Extern::Memory(m)) => m,
_ => {
return Err(wasmtime::Trap::new("missing required memory export"));
}
};
let mem = #runtime::WasmtimeGuestMemory::new(mem);
match #target_module::#name_ident(ctx, &mem #(, #arg_names)*) #await_ {
Ok(r) => Ok(r.into()),
Err(wasmtime_wiggle::Trap::String(err)) => Err(wasmtime::Trap::new(err)),
Err(wasmtime_wiggle::Trap::I32Exit(err)) => Err(wasmtime::Trap::i32_exit(err)),
}
}
}
});
match asyncness {
Asyncness::Async => {
let wrapper = format_ident!("wrap{}_async", params.len());
ctors.push(quote! {
let #name_ident = wasmtime::Func::#wrapper(
store,
ctx.clone(),
move |caller: wasmtime::Caller<'_>, my_ctx: &std::rc::Rc<std::cell::RefCell<_>> #(,#arg_decls)*|
-> Box<dyn std::future::Future<Output = Result<#ret_ty, wasmtime::Trap>>> {
Box::new(async move { Self::#fn_ident(&caller, &mut my_ctx.borrow_mut() #(, #arg_names)*).await })
}
);
});
}
Asyncness::Blocking => {
ctors.push(quote! {
let my_ctx = ctx.clone();
let #name_ident = wasmtime::Func::wrap(
store,
move |caller: wasmtime::Caller #(, #arg_decls)*| -> Result<#ret_ty, wasmtime::Trap> {
#rt::run_in_dummy_executor(Self::#fn_ident(&caller, &mut my_ctx.borrow_mut() #(, #arg_names)*))
}
);
});
}
Asyncness::Sync => {
ctors.push(quote! {
let my_ctx = ctx.clone();
let #name_ident = wasmtime::Func::wrap(
store,
move |caller: wasmtime::Caller #(, #arg_decls)*| -> Result<#ret_ty, wasmtime::Trap> {
Self::#fn_ident(&caller, &mut my_ctx.borrow_mut() #(, #arg_names)*)
}
);
});
}
}
let host_wrapper = match asyncness {
Asyncness::Async => {
let wrapper = format_ident!("wrap{}_host_func_async", params.len());
quote! {
config.#wrapper(
module,
field,
move |caller #(,#arg_decls)*|
-> Box<dyn std::future::Future<Output = Result<#ret_ty, wasmtime::Trap>>> {
Box::new(async move {
let ctx = caller.store()
.get::<std::rc::Rc<std::cell::RefCell<#ctx_type>>>()
.ok_or_else(|| wasmtime::Trap::new("context is missing in the store"))?;
let result = Self::#fn_ident(&caller, &mut ctx.borrow_mut() #(, #arg_names)*).await;
result
})
}
);
}
}
Asyncness::Blocking => {
quote! {
config.wrap_host_func(
module,
field,
move |caller: wasmtime::Caller #(, #arg_decls)*| -> Result<#ret_ty, wasmtime::Trap> {
let ctx = caller
.store()
.get::<std::rc::Rc<std::cell::RefCell<#ctx_type>>>()
.ok_or_else(|| wasmtime::Trap::new("context is missing in the store"))?;
#rt::run_in_dummy_executor(Self::#fn_ident(&caller, &mut ctx.borrow_mut() #(, #arg_names)*))
},
);
}
}
Asyncness::Sync => {
quote! {
config.wrap_host_func(
module,
field,
move |caller: wasmtime::Caller #(, #arg_decls)*| -> Result<#ret_ty, wasmtime::Trap> {
let ctx = caller
.store()
.get::<std::rc::Rc<std::cell::RefCell<#ctx_type>>>()
.ok_or_else(|| wasmtime::Trap::new("context is missing in the store"))?;
Self::#fn_ident(&caller, &mut ctx.borrow_mut() #(, #arg_names)*)
},
);
}
}
};
host_funcs.push((func.name.clone(), host_wrapper));
}