mate_task/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{FnArg, ItemFn, ReturnType, Type, parse_macro_input};
4
5#[proc_macro_attribute]
6pub fn mate_object(_attr: TokenStream, item: TokenStream) -> TokenStream {
7    let input = parse_macro_input!(item as syn::DeriveInput);
8    let name = &input.ident;
9
10    let expanded = quote! {
11        #[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
12        #input
13
14        const _: () = {
15            fn assert_impl<T: serde::Serialize + serde::de::DeserializeOwned + Clone + std::fmt::Debug>() {}
16            fn assert_type() {
17                assert_impl::<#name>();
18            }
19        };
20    };
21
22    TokenStream::from(expanded)
23}
24
25#[proc_macro_attribute]
26pub fn mate_handler(_attr: TokenStream, item: TokenStream) -> TokenStream {
27    let input_fn = parse_macro_input!(item as ItemFn);
28    let fn_name = &input_fn.sig.ident;
29    let fn_vis = &input_fn.vis;
30    let fn_block = &input_fn.block;
31    let fn_attrs = &input_fn.attrs;
32    let fn_sig = &input_fn.sig;
33
34    // Extract the input type from the first parameter
35    let input_type = match fn_sig.inputs.first() {
36        Some(FnArg::Typed(pat_type)) => &pat_type.ty,
37        _ => {
38            return syn::Error::new_spanned(
39                &fn_sig.inputs,
40                "Handler function must have at least one parameter",
41            )
42            .to_compile_error()
43            .into();
44        }
45    };
46
47    // Extract the output type from the return type
48    let output_type = match &fn_sig.output {
49        ReturnType::Type(_, ty) => {
50            // Handle Result<T, E> -> extract T
51            if let Type::Path(type_path) = ty.as_ref() {
52                if let Some(segment) = type_path.path.segments.last() {
53                    if segment.ident == "Result" {
54                        // Extract the Ok type from Result<T, E>
55                        if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
56                            if let Some(syn::GenericArgument::Type(ok_type)) = args.args.first() {
57                                ok_type
58                            } else {
59                                ty.as_ref()
60                            }
61                        } else {
62                            ty.as_ref()
63                        }
64                    } else {
65                        ty.as_ref()
66                    }
67                } else {
68                    ty.as_ref()
69                }
70            } else {
71                ty.as_ref()
72            }
73        }
74        ReturnType::Default => {
75            return syn::Error::new_spanned(
76                &fn_sig.output,
77                "Handler function must return a Result type",
78            )
79            .to_compile_error()
80            .into();
81        }
82    };
83
84    let expanded = quote! {
85        use anyhow::Result;
86        use wit_bindgen;
87
88        mod bindings {
89            wit_bindgen::generate!({
90                async: true,
91                inline: r"
92                    package mate:runtime@0.1.0;
93
94                    world mate {
95                        export handler: async func(data: string) -> result<string, string>;
96                    }",
97            });
98        }
99
100        struct Mate;
101
102        impl bindings::Guest for Mate {
103            async fn handler(data: String) -> Result<String, String> {
104                let input: #input_type = serde_json::from_str(&data)
105                    .map_err(|e| format!("Failed to deserialize input: {}", e))?;
106
107                let result: Result<#output_type> = #fn_name(input).await;
108
109                match result {
110                    Ok(val) => serde_json::to_string(&val)
111                        .map_err(|e| format!("Failed to serialize output: {}", e)),
112                    Err(err) => Err(format!("Handler error: {}", err)),
113                }
114            }
115        }
116
117        bindings::export!(Mate with_types_in bindings);
118
119        #(#fn_attrs)*
120        #fn_vis #fn_sig {
121            #fn_block
122        }
123    };
124
125    TokenStream::from(expanded)
126}
mate_task/lib.rs

mate_task/
lib.rs
