wiggle_generate/
wasmtime.rs

1use crate::CodegenSettings;
2use crate::config::Asyncness;
3use crate::funcs::func_bounds;
4use crate::names;
5use proc_macro2::{Ident, Span, TokenStream};
6use quote::{format_ident, quote};
7use std::collections::HashSet;
8
9pub fn link_module(
10    module: &witx::Module,
11    target_path: Option<&syn::Path>,
12    settings: &CodegenSettings,
13) -> TokenStream {
14    let module_ident = names::module(&module.name);
15
16    let send_bound = if settings.async_.contains_async(module) {
17        quote! { + Send, T: Send }
18    } else {
19        quote! {}
20    };
21
22    let mut bodies = Vec::new();
23    let mut bounds = HashSet::new();
24    for f in module.funcs() {
25        let asyncness = settings.async_.get(module.name.as_str(), f.name.as_str());
26        bodies.push(generate_func(&module, &f, target_path, asyncness));
27        let bound = func_bounds(module, &f, settings);
28        for b in bound {
29            bounds.insert(b);
30        }
31    }
32
33    let ctx_bound = if let Some(target_path) = target_path {
34        let bounds = bounds
35            .into_iter()
36            .map(|b| quote!(#target_path::#module_ident::#b));
37        quote!( #(#bounds)+* #send_bound )
38    } else {
39        let bounds = bounds.into_iter();
40        quote!( #(#bounds)+* #send_bound )
41    };
42
43    let func_name = if target_path.is_none() {
44        format_ident!("add_to_linker")
45    } else {
46        format_ident!("add_{}_to_linker", module_ident)
47    };
48
49    let u = if settings.mutable {
50        quote!(&mut U)
51    } else {
52        quote!(&U)
53    };
54    quote! {
55        /// Adds all instance items to the specified `Linker`.
56        pub fn #func_name<T, U>(
57            linker: &mut wiggle::wasmtime_crate::Linker<T>,
58            get_cx: impl Fn(&mut T) -> #u + Send + Sync + Copy + 'static,
59        ) -> wiggle::anyhow::Result<()>
60            where
61                T: 'static,
62                U: #ctx_bound #send_bound
63        {
64            #(#bodies)*
65            Ok(())
66        }
67    }
68}
69
70fn generate_func(
71    module: &witx::Module,
72    func: &witx::InterfaceFunc,
73    target_path: Option<&syn::Path>,
74    asyncness: Asyncness,
75) -> TokenStream {
76    let module_str = module.name.as_str();
77    let module_ident = names::module(&module.name);
78
79    let field_str = func.name.as_str();
80    let field_ident = names::func(&func.name);
81
82    let (params, results) = func.wasm_signature();
83
84    let arg_names = (0..params.len())
85        .map(|i| Ident::new(&format!("arg{i}"), Span::call_site()))
86        .collect::<Vec<_>>();
87    let arg_tys = params
88        .iter()
89        .map(|ty| names::wasm_type(*ty))
90        .collect::<Vec<_>>();
91    let arg_decls = arg_names
92        .iter()
93        .zip(arg_tys.iter())
94        .map(|(name, ty)| {
95            quote! { #name: #ty }
96        })
97        .collect::<Vec<_>>();
98
99    let ret_ty = match results.len() {
100        0 => quote!(()),
101        1 => names::wasm_type(results[0]),
102        _ => unimplemented!(),
103    };
104
105    let await_ = if asyncness.is_sync() {
106        quote!()
107    } else {
108        quote!(.await)
109    };
110
111    let abi_func = if let Some(target_path) = target_path {
112        quote!( #target_path::#module_ident::#field_ident )
113    } else {
114        quote!( #field_ident )
115    };
116
117    let body = quote! {
118        let export = caller.get_export("memory");
119        let (mut mem, ctx) = match &export {
120            Some(wiggle::wasmtime_crate::Extern::Memory(m)) => {
121                let (mem, ctx) = m.data_and_store_mut(&mut caller);
122                let ctx = get_cx(ctx);
123                (wiggle::GuestMemory::Unshared(mem), ctx)
124            }
125            Some(wiggle::wasmtime_crate::Extern::SharedMemory(m)) => {
126                let ctx = get_cx(caller.data_mut());
127                (wiggle::GuestMemory::Shared(m.data()), ctx)
128            }
129            _ => wiggle::anyhow::bail!("missing required memory export"),
130        };
131        Ok(<#ret_ty>::from(#abi_func(ctx, &mut mem #(, #arg_names)*) #await_ ?))
132    };
133
134    match asyncness {
135        Asyncness::Async => {
136            let arg_decls = quote! { ( #(#arg_names,)* ) : ( #(#arg_tys,)* ) };
137            quote! {
138                linker.func_wrap_async(
139                    #module_str,
140                    #field_str,
141                    move |mut caller: wiggle::wasmtime_crate::Caller<'_, T>, #arg_decls| {
142                        Box::new(async move { #body })
143                    },
144                )?;
145            }
146        }
147
148        Asyncness::Blocking { block_with } => {
149            quote! {
150                linker.func_wrap(
151                    #module_str,
152                    #field_str,
153                    move |mut caller: wiggle::wasmtime_crate::Caller<'_, T> #(, #arg_decls)*| -> wiggle::anyhow::Result<#ret_ty> {
154                        let result = async { #body };
155                        #block_with(result)?
156                    },
157                )?;
158            }
159        }
160
161        Asyncness::Sync => {
162            quote! {
163                linker.func_wrap(
164                    #module_str,
165                    #field_str,
166                    move |mut caller: wiggle::wasmtime_crate::Caller<'_, T> #(, #arg_decls)*| -> wiggle::anyhow::Result<#ret_ty> {
167                        #body
168                    },
169                )?;
170            }
171        }
172    }
173}