wiggle_generate/
wasmtime.rs1use 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 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}