worker_codegen/
wit.rs

1use convert_case::{Case, Casing};
2use proc_macro2::TokenStream;
3use quote::ToTokens;
4use quote::{format_ident, quote};
5use syn::Ident;
6use wit_parser::Interface;
7
8fn path_type(name: &str) -> anyhow::Result<syn::Type> {
9    let ty: syn::TypePath = syn::parse_str(name)?;
10    Ok(syn::Type::Path(ty))
11}
12
13fn wit_type_to_syn(ty: &wit_parser::Type) -> anyhow::Result<syn::Type> {
14    path_type(&wit_type_to_str(ty)?)
15}
16
17fn wit_type_to_str(ty: &wit_parser::Type) -> anyhow::Result<String> {
18    Ok(match ty {
19        wit_parser::Type::Bool => "bool".to_string(),
20        wit_parser::Type::U8 => "u8".to_string(),
21        wit_parser::Type::U16 => "u16".to_string(),
22        wit_parser::Type::U32 => "u32".to_string(),
23        wit_parser::Type::U64 => "u64".to_string(),
24        wit_parser::Type::S8 => "i8".to_string(),
25        wit_parser::Type::S16 => "i16".to_string(),
26        wit_parser::Type::S32 => "i32".to_string(),
27        wit_parser::Type::S64 => "i64".to_string(),
28        wit_parser::Type::F32 => "f32".to_string(),
29        wit_parser::Type::F64 => "f64".to_string(),
30        wit_parser::Type::Char => "char".to_string(),
31        wit_parser::Type::String => "String".to_string(),
32        wit_parser::Type::Id(t) => anyhow::bail!("Unsupported type: '{:?}'", t),
33    })
34}
35
36fn expand_args(method: &wit_parser::Function) -> anyhow::Result<Vec<syn::FnArg>> {
37    let mut args = Vec::with_capacity(method.params.len());
38    for (arg_name, arg) in &method.params {
39        let param = syn::FnArg::Typed(syn::PatType {
40            attrs: vec![],
41            pat: Box::new(syn::Pat::Ident(syn::PatIdent {
42                attrs: vec![],
43                by_ref: None,
44                mutability: None,
45                ident: format_ident!("{}", arg_name),
46                subpat: None,
47            })),
48            colon_token: Default::default(),
49            ty: Box::new(wit_type_to_syn(arg)?),
50        });
51        args.push(param);
52    }
53    Ok(args)
54}
55
56fn expand_trait(interface: &Interface, interface_name: &Ident) -> anyhow::Result<syn::ItemTrait> {
57    let trait_raw = quote!(
58        #[async_trait::async_trait]
59        pub trait #interface_name {
60        }
61    );
62    let mut trait_item: syn::ItemTrait = syn::parse2(trait_raw)?;
63
64    for (name, method) in &interface.functions {
65        let ident = format_ident!("{}", name.to_case(Case::Snake));
66        let ret_type = if let wit_parser::Results::Anon(ty) = &method.results {
67            format_ident!("{}", wit_type_to_str(ty)?)
68        } else {
69            anyhow::bail!("Unsupported return type: '{:?}'", method.results);
70        };
71
72        let method_raw = quote!(
73            // TODO: docs
74            async fn #ident(&self) -> ::worker::Result<#ret_type>;
75        );
76
77        let mut method_item: syn::TraitItemFn = syn::parse2(method_raw)?;
78
79        method_item.sig.inputs.extend(expand_args(method)?);
80        trait_item.items.push(syn::TraitItem::Fn(method_item));
81    }
82
83    Ok(trait_item)
84}
85
86fn expand_struct(struct_name: &Ident, sys_name: &Ident) -> anyhow::Result<syn::ItemStruct> {
87    let struct_raw = quote!(
88        pub struct #struct_name(::worker::send::SendWrapper<sys::#sys_name>);
89    );
90    let struct_item: syn::ItemStruct = syn::parse2(struct_raw)?;
91    Ok(struct_item)
92}
93
94fn expand_from_impl(struct_name: &Ident, from_type: &syn::Type) -> anyhow::Result<syn::ItemImpl> {
95    let impl_raw = quote!(
96        impl From<#from_type> for #struct_name {
97            fn from(fetcher: #from_type) -> Self {
98                Self(::worker::send::SendWrapper::new(fetcher.into_rpc()))
99            }
100        }
101    );
102    let impl_item: syn::ItemImpl = syn::parse2(impl_raw)?;
103    Ok(impl_item)
104}
105
106fn expand_rpc_impl(
107    interface: &Interface,
108    interface_name: &Ident,
109    struct_name: &Ident,
110) -> anyhow::Result<syn::ItemImpl> {
111    let impl_raw = quote!(
112        #[async_trait::async_trait]
113        impl #interface_name for #struct_name {}
114    );
115    let mut impl_item: syn::ItemImpl = syn::parse2(impl_raw)?;
116
117    for (name, method) in &interface.functions {
118        println!("\tFound method: '{}'.", name);
119        let ident = format_ident!("{}", name.to_case(Case::Snake));
120        let invocation_raw = quote!(self.0.#ident());
121        let mut invocation_item: syn::ExprMethodCall = syn::parse2(invocation_raw)?;
122        for (arg_name, _) in &method.params {
123            let mut segments = syn::punctuated::Punctuated::new();
124            segments.push(syn::PathSegment {
125                ident: format_ident!("{}", arg_name),
126                arguments: syn::PathArguments::None,
127            });
128            invocation_item.args.push(syn::Expr::Path(syn::ExprPath {
129                attrs: vec![],
130                qself: None,
131                path: syn::Path {
132                    leading_colon: None,
133                    segments,
134                },
135            }));
136        }
137
138        let ret_type = if let wit_parser::Results::Anon(ty) = &method.results {
139            format_ident!("{}", wit_type_to_str(ty)?)
140        } else {
141            anyhow::bail!("Unsupported return type: '{:?}'", method.results);
142        };
143
144        let method_raw = quote!(
145            async fn #ident(&self) -> ::worker::Result<#ret_type> {
146                let promise = #invocation_item?;
147                let fut = ::worker::send::SendFuture::new(::worker::wasm_bindgen_futures::JsFuture::from(promise));
148                let output = fut.await?;
149                Ok(::serde_wasm_bindgen::from_value(output)?)
150            }
151        );
152
153        let mut method_item: syn::ImplItemFn = syn::parse2(method_raw)?;
154        method_item.sig.inputs.extend(expand_args(method)?);
155        impl_item.items.push(syn::ImplItem::Fn(method_item));
156    }
157    Ok(impl_item)
158}
159
160fn expand_sys_module(interface: &Interface, sys_name: &Ident) -> anyhow::Result<syn::ItemMod> {
161    let f_mod_raw = quote!(
162        #[wasm_bindgen]
163        extern "C" {
164            #[wasm_bindgen(extends=::worker::js_sys::Object)]
165            pub type #sys_name;
166        }
167    );
168    let mut f_mod_item: syn::ItemForeignMod = syn::parse2(f_mod_raw)?;
169
170    for (name, method) in &interface.functions {
171        let ident = format_ident!("{}", name.to_case(Case::Snake));
172        let extern_name = name.to_case(Case::Camel);
173        let method_raw = quote!(
174            #[wasm_bindgen(method, catch, js_name = #extern_name)]
175            // TODO: args
176            pub fn #ident(
177                this: &#sys_name,
178            ) -> std::result::Result<::worker::js_sys::Promise, ::worker::wasm_bindgen::JsValue>;
179        );
180        let mut method_item: syn::ForeignItemFn = syn::parse2(method_raw)?;
181        method_item.sig.inputs.extend(expand_args(method)?);
182        f_mod_item.items.push(syn::ForeignItem::Fn(method_item));
183    }
184
185    let mod_raw = quote!(
186        mod sys {
187            use ::wasm_bindgen::prelude::*;
188        }
189    );
190    let mut mod_item: syn::ItemMod = syn::parse2(mod_raw)?;
191    if let Some(ref mut content) = mod_item.content {
192        content.1.push(syn::Item::ForeignMod(f_mod_item));
193    }
194
195    Ok(mod_item)
196}
197
198fn expand_wit(path: &str) -> anyhow::Result<syn::File> {
199    let mut resolver = wit_parser::Resolve::new();
200    resolver.push_file(path)?;
201
202    // Items: Trait, Struct, Trait Impl, From Impl, Sys Module
203    let mut items = Vec::with_capacity(resolver.interfaces.len() * 5);
204
205    for (_, interface) in resolver.interfaces {
206        let name = interface.name.clone().unwrap();
207        println!("Found Interface: '{}'", name);
208        let interface_name = format_ident!("{}", name.to_case(Case::Pascal));
209        println!("Generating Trait '{}'", interface_name);
210        let struct_name = format_ident!("{}Service", interface_name);
211        let sys_name = format_ident!("{}Sys", interface_name);
212
213        // Sys Module
214        items.push(syn::Item::Mod(expand_sys_module(&interface, &sys_name)?));
215        //  Trait
216        items.push(syn::Item::Trait(expand_trait(&interface, &interface_name)?));
217        // Struct
218        items.push(syn::Item::Struct(expand_struct(&struct_name, &sys_name)?));
219        // Trait Impl
220        items.push(syn::Item::Impl(expand_rpc_impl(
221            &interface,
222            &interface_name,
223            &struct_name,
224        )?));
225        // From Impl for Fetcher and Stub
226        items.push(syn::Item::Impl(expand_from_impl(
227            &struct_name,
228            &syn::parse_str("::worker::Fetcher")?,
229        )?));
230        items.push(syn::Item::Impl(expand_from_impl(
231            &struct_name,
232            &syn::parse_str("::worker::Stub")?,
233        )?));
234    }
235
236    let rust_file = syn::File {
237        shebang: None,
238        attrs: vec![],
239        items,
240    };
241    Ok(rust_file)
242}
243
244/// Expands a WIT file into a Rust source file as a string.
245pub fn expand_wit_source(path: &str) -> anyhow::Result<String> {
246    let file = expand_wit(path)?;
247    Ok(prettyplease::unparse(&file))
248}
249
250/// Expands a WIT file into a Rust source file as a token stream.
251pub fn expand_wit_tokens(path: &str) -> anyhow::Result<TokenStream> {
252    let file = expand_wit(path)?;
253    Ok(file.into_token_stream())
254}