1use proc_macro::TokenStream;
9use quote::{format_ident, quote};
10
11fn schema_derive() -> proc_macro2::TokenStream {
16 quote! {
17 #[derive(::yog_api::yog_data::Schema)]
18 #[data(crate = "::yog_api::yog_data")]
19 }
20}
21
22fn is_registry_ty(ty: &syn::Type) -> bool {
25 let ty = match ty {
26 syn::Type::Reference(r) => &*r.elem,
27 other => other,
28 };
29 quote! { #ty }
30 .to_string()
31 .replace(' ', "")
32 .ends_with("Registry")
33}
34
35fn call_params(sig: &syn::Signature) -> (bool, Vec<(syn::Ident, syn::Type)>) {
38 let mut typed: Vec<&syn::PatType> = sig
39 .inputs
40 .iter()
41 .filter_map(|a| match a {
42 syn::FnArg::Typed(pt) => Some(pt),
43 syn::FnArg::Receiver(_) => None,
44 })
45 .collect();
46 let registry = typed.first().is_some_and(|pt| is_registry_ty(&pt.ty));
47 if registry {
48 typed.remove(0);
49 }
50 let params = typed
51 .iter()
52 .enumerate()
53 .map(|(i, pt)| {
54 let name = match &*pt.pat {
55 syn::Pat::Ident(p) => p.ident.clone(),
56 _ => format_ident!("__yog_arg{}", i),
57 };
58 (name, (*pt.ty).clone())
59 })
60 .collect();
61 (registry, params)
62}
63
64fn args_struct(
66 name: &syn::Ident,
67 params: &[(syn::Ident, syn::Type)],
68) -> (syn::Ident, proc_macro2::TokenStream) {
69 let args = format_ident!("__YogArgs_{}", name);
70 let fields = params.iter().map(|(n, t)| quote! { #n: #t });
71 let derive = schema_derive();
72 (
73 args.clone(),
74 quote! {
75 #[doc(hidden)]
76 #[allow(non_camel_case_types)]
78 #derive
79 struct #args { #(#fields),* }
80 },
81 )
82}
83
84fn output_type(sig: &syn::Signature) -> proc_macro2::TokenStream {
85 match &sig.output {
86 syn::ReturnType::Default => quote! { () },
87 syn::ReturnType::Type(_, ty) => quote! { #ty },
88 }
89}
90
91fn doc_of(attrs: &[syn::Attribute]) -> String {
95 let lines: Vec<String> = attrs
96 .iter()
97 .filter(|a| a.path().is_ident("doc"))
98 .filter_map(|a| match &a.meta {
99 syn::Meta::NameValue(nv) => match &nv.value {
100 syn::Expr::Lit(syn::ExprLit {
101 lit: syn::Lit::Str(s),
102 ..
103 }) => Some(s.value().trim().to_owned()),
104 _ => None,
105 },
106 _ => None,
107 })
108 .collect();
109 lines.join(" ")
110}
111
112fn implements_of(attr: TokenStream) -> syn::Result<Vec<String>> {
114 let mut out = Vec::new();
115 let parser = syn::meta::parser(|meta| {
116 if meta.path.is_ident("implements") {
117 out.push(meta.value()?.parse::<syn::LitStr>()?.value());
118 Ok(())
119 } else {
120 Err(meta.error("unknown `yog_export` argument (implements)"))
121 }
122 });
123 syn::parse::Parser::parse(parser, attr)?;
124 Ok(out)
125}
126
127#[proc_macro_attribute]
128pub fn yog_export(attr: TokenStream, item: TokenStream) -> TokenStream {
129 let implements = match implements_of(attr) {
130 Ok(i) => i,
131 Err(e) => return e.to_compile_error().into(),
132 };
133 let tokens = proc_macro2::TokenStream::from(item.clone());
134 if syn::parse2::<syn::ItemStruct>(tokens.clone()).is_ok()
135 || syn::parse2::<syn::ItemEnum>(tokens.clone()).is_ok()
136 {
137 let derive = schema_derive();
138 return TokenStream::from(quote! { #derive #tokens });
139 }
140 let func = syn::parse_macro_input!(item as syn::ItemFn);
141 let name = &func.sig.ident;
142 let name_str = name.to_string();
143 let serve_name = format_ident!("__yog_serve_{}", name);
144 let ctor_name = format_ident!("__yog_ctor_register_{}", name);
145 let (registry, params) = call_params(&func.sig);
146 let (args, args_def) = args_struct(name, ¶ms);
147 let output = output_type(&func.sig);
148 let names: Vec<&syn::Ident> = params.iter().map(|(n, _)| n).collect();
149 let doc = doc_of(&func.attrs);
150 let (setup, registry_arg) = if registry {
152 (
153 quote! { let mut __yog_registry = unsafe { ::yog_api::Registry::from_raw(::yog_api::__current_api_ptr() as *const ::yog_api::YogApi) }; },
154 quote! { &mut __yog_registry, },
155 )
156 } else {
157 (quote! {}, quote! {})
158 };
159 TokenStream::from(quote! {
160 #func
161
162 #args_def
163
164 #[doc(hidden)]
165 unsafe extern "C" fn #serve_name(
166 _ud: *mut ::std::os::raw::c_void,
167 payload: ::yog_api::YogBytes,
168 reply: *const ::yog_api::YogReply,
169 ) -> ::yog_api::YogStatus {
170 unsafe {
171 ::yog_api::__yog_serve::<#args, #output>(payload, reply, |a| {
172 #setup
173 #name(#registry_arg #(a.#names),*)
174 })
175 }
176 }
177
178 #[::yog_api::ctor::ctor(unsafe, crate_path = ::yog_api::ctor)]
180 #[doc(hidden)]
181 fn #ctor_name() {
182 ::yog_api::__yog_export_registry().lock().push(
183 ::yog_api::YogExportEntry {
184 name: #name_str,
185 handler: #serve_name,
186 doc: #doc,
187 shapes: || (<#args as ::yog_api::yog_data::Schema>::shape(), <#output as ::yog_api::yog_data::Schema>::shape()),
188 implements: &[#(#implements),*],
189 }
190 );
191 }
192 })
193}
194
195#[proc_macro_attribute]
202pub fn bench(_attr: TokenStream, item: TokenStream) -> TokenStream {
203 let mut f = syn::parse_macro_input!(item as syn::ItemFn);
204 if matches!(f.vis, syn::Visibility::Inherited) {
205 f.vis = syn::parse_quote!(pub(crate));
206 }
207 TokenStream::from(quote! { #f })
208}
209
210#[proc_macro]
217pub fn import(input: TokenStream) -> TokenStream {
218 let import_block = syn::parse_macro_input!(input as ImportBlock);
219 let mod_name = &import_block.mod_name;
220 let mod_ident = format_ident!("{}", mod_name.replace('-', "_"));
221 let mut items = proc_macro2::TokenStream::new();
222 for item in &import_block.items {
223 match item {
224 ImportItem::Func(func) => items.extend(generate_import_fn(func, mod_name)),
225 ImportItem::Type(t) => {
226 let derive = schema_derive();
227 items.extend(quote! { #derive #t });
228 }
229 }
230 }
231 TokenStream::from(quote! { pub mod #mod_ident { use super::*; #items } })
232}
233
234struct ImportFn {
236 vis: syn::Visibility,
237 sig: syn::Signature,
238}
239
240fn generate_import_fn(func: &ImportFn, mod_name: &str) -> proc_macro2::TokenStream {
241 let name = &func.sig.ident;
242 let name_str = name.to_string();
243 let vis = &func.vis;
244 let (_, params) = call_params(&func.sig);
245 let (args, args_def) = args_struct(name, ¶ms);
246 let output = output_type(&func.sig);
247 let names: Vec<&syn::Ident> = params.iter().map(|(n, _)| n).collect();
248 let types: Vec<&syn::Type> = params.iter().map(|(_, t)| t).collect();
249 quote! {
250 #args_def
251
252 #vis fn #name(#(#names: #types),*) -> ::std::option::Option<#output> {
254 ::yog_api::call_export::<#args, #output>(#mod_name, #name_str, &#args { #(#names),* })
255 }
256 }
257}
258
259struct ImportBlock {
262 mod_name: String,
263 items: Vec<ImportItem>,
264}
265enum ImportItem {
266 Func(Box<ImportFn>),
267 Type(proc_macro2::TokenStream),
269}
270
271impl syn::parse::Parse for ImportBlock {
272 fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
273 let from_kw: syn::Ident = input.parse()?;
274 if from_kw != "from" {
275 return Err(syn::Error::new(from_kw.span(), "expected `from`"));
276 }
277 let lit: syn::LitStr = input.parse()?;
278 let mod_name = lit.value();
279 let content;
280 syn::braced!(content in input);
281 let mut items = Vec::new();
282 while !content.is_empty() {
283 let attrs = content.call(syn::Attribute::parse_outer)?;
284 let vis: syn::Visibility = content.parse()?;
285 if content.peek(syn::Token![fn]) {
286 let sig: syn::Signature = content.parse()?;
287 if content.peek(syn::token::Brace) {
288 let _body: syn::Block = content.parse()?;
289 } else {
290 content.parse::<syn::Token![;]>()?;
291 }
292 items.push(ImportItem::Func(Box::new(ImportFn { vis, sig })));
293 } else if content.peek(syn::Token![enum]) {
294 let mut e: syn::ItemEnum = content.parse()?;
295 e.attrs = attrs;
296 e.vis = vis;
297 items.push(ImportItem::Type(quote! { #e }));
298 } else {
299 let mut s: syn::ItemStruct = content.parse()?;
300 s.attrs = attrs;
301 s.vis = vis;
302 items.push(ImportItem::Type(quote! { #s }));
303 let _ = content.parse::<syn::Token![;]>();
304 }
305 }
306 Ok(ImportBlock { mod_name, items })
307 }
308}