wiggle_generate/
wasmtime.rs

1use crate::config::Asyncness;
2use crate::funcs::func_bounds;
3use crate::names;
4use crate::CodegenSettings;
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                U: #ctx_bound #send_bound
62        {
63            #(#bodies)*
64            Ok(())
65        }
66    }
67}
68
69fn generate_func(
70    module: &witx::Module,
71    func: &witx::InterfaceFunc,
72    target_path: Option<&syn::Path>,
73    asyncness: Asyncness,
74) -> TokenStream {
75    let module_str = module.name.as_str();
76    let module_ident = names::module(&module.name);
77
78    let field_str = func.name.as_str();
79    let field_ident = names::func(&func.name);
80
81    let (params, results) = func.wasm_signature();
82
83    let arg_names = (0..params.len())
84        .map(|i| Ident::new(&format!("arg{i}"), Span::call_site()))
85        .collect::<Vec<_>>();
86    let arg_tys = params
87        .iter()
88        .map(|ty| names::wasm_type(*ty))
89        .collect::<Vec<_>>();
90    let arg_decls = arg_names
91        .iter()
92        .zip(arg_tys.iter())
93        .map(|(name, ty)| {
94            quote! { #name: #ty }
95        })
96        .collect::<Vec<_>>();
97
98    let ret_ty = match results.len() {
99        0 => quote!(()),
100        1 => names::wasm_type(results[0]),
101        _ => unimplemented!(),
102    };
103
104    let await_ = if asyncness.is_sync() {
105        quote!()
106    } else {
107        quote!(.await)
108    };
109
110    let abi_func = if let Some(target_path) = target_path {
111        quote!( #target_path::#module_ident::#field_ident )
112    } else {
113        quote!( #field_ident )
114    };
115
116    let body = quote! {
117        let export = caller.get_export("memory");
118        let (mut mem, ctx) = match &export {
119            Some(wiggle::wasmtime_crate::Extern::Memory(m)) => {
120                let (mem, ctx) = m.data_and_store_mut(&mut caller);
121                let ctx = get_cx(ctx);
122                (wiggle::GuestMemory::Unshared(mem), ctx)
123            }
124            Some(wiggle::wasmtime_crate::Extern::SharedMemory(m)) => {
125                let ctx = get_cx(caller.data_mut());
126                (wiggle::GuestMemory::Shared(m.data()), ctx)
127            }
128            _ => wiggle::anyhow::bail!("missing required memory export"),
129        };
130        Ok(<#ret_ty>::from(#abi_func(ctx, &mut mem #(, #arg_names)*) #await_ ?))
131    };
132
133    match asyncness {
134        Asyncness::Async => {
135            let arg_decls = quote! { ( #(#arg_names,)* ) : ( #(#arg_tys,)* ) };
136            quote! {
137                linker.func_wrap_async(
138                    #module_str,
139                    #field_str,
140                    move |mut caller: wiggle::wasmtime_crate::Caller<'_, T>, #arg_decls| {
141                        Box::new(async move { #body })
142                    },
143                )?;
144            }
145        }
146
147        Asyncness::Blocking { block_with } => {
148            quote! {
149                linker.func_wrap(
150                    #module_str,
151                    #field_str,
152                    move |mut caller: wiggle::wasmtime_crate::Caller<'_, T> #(, #arg_decls)*| -> wiggle::anyhow::Result<#ret_ty> {
153                        let result = async { #body };
154                        #block_with(result)?
155                    },
156                )?;
157            }
158        }
159
160        Asyncness::Sync => {
161            quote! {
162                linker.func_wrap(
163                    #module_str,
164                    #field_str,
165                    move |mut caller: wiggle::wasmtime_crate::Caller<'_, T> #(, #arg_decls)*| -> wiggle::anyhow::Result<#ret_ty> {
166                        #body
167                    },
168                )?;
169            }
170        }
171    }
172}