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