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 enum_tag(e: &syn::ItemEnum) -> proc_macro2::TokenStream {
25 use yog_ids::data_attr::{DATA, TAG};
26 let has_data = e
27 .variants
28 .iter()
29 .any(|v| !matches!(v.fields, syn::Fields::Unit));
30 let tagged = e.attrs.iter().filter(|a| a.path().is_ident(DATA)).any(|a| {
31 let mut tag = false;
32 let _ = a.parse_nested_meta(|m| {
34 tag |= m.path.is_ident(TAG);
35 if m.input.peek(syn::Token![=]) {
36 m.value()?.parse::<syn::Expr>()?;
37 }
38 Ok(())
39 });
40 tag
41 });
42 if !has_data || tagged {
43 return quote! {};
44 }
45 let tag = yog_ids::word::TYPE;
46 let (data, tag_key) = (format_ident!("{DATA}"), format_ident!("{TAG}"));
47 quote! { #[#data(#tag_key = #tag)] }
48}
49
50fn is_registry_ty(ty: &syn::Type) -> bool {
53 let ty = match ty {
54 syn::Type::Reference(r) => &*r.elem,
55 other => other,
56 };
57 quote! { #ty }
58 .to_string()
59 .replace(' ', "")
60 .ends_with("Registry")
61}
62
63fn call_params(sig: &syn::Signature) -> (bool, Vec<(syn::Ident, syn::Type)>) {
66 let mut typed: Vec<&syn::PatType> = sig
67 .inputs
68 .iter()
69 .filter_map(|a| match a {
70 syn::FnArg::Typed(pt) => Some(pt),
71 syn::FnArg::Receiver(_) => None,
72 })
73 .collect();
74 let registry = typed.first().is_some_and(|pt| is_registry_ty(&pt.ty));
75 if registry {
76 typed.remove(0);
77 }
78 let params = typed
79 .iter()
80 .enumerate()
81 .map(|(i, pt)| {
82 let name = match &*pt.pat {
83 syn::Pat::Ident(p) => p.ident.clone(),
84 _ => format_ident!("__yog_arg{}", i),
85 };
86 (name, (*pt.ty).clone())
87 })
88 .collect();
89 (registry, params)
90}
91
92struct Args {
96 ty: proc_macro2::TokenStream,
98 def: proc_macro2::TokenStream,
100 pack: proc_macro2::TokenStream,
102 unpack: proc_macro2::TokenStream,
104}
105
106fn args_of(name: &syn::Ident, params: &[(syn::Ident, syn::Type)]) -> Args {
107 if let [(param, ty)] = params {
108 return Args {
109 ty: quote! { #ty },
110 def: quote! {},
111 pack: quote! { &#param },
112 unpack: quote! { a },
113 };
114 }
115 let args = format_ident!("__YogArgs_{}", name);
116 let fields = params.iter().map(|(n, t)| quote! { #n: #t });
117 let names: Vec<&syn::Ident> = params.iter().map(|(n, _)| n).collect();
118 let derive = schema_derive();
119 Args {
120 ty: quote! { #args },
121 def: quote! {
122 #[doc(hidden)]
123 #[allow(non_camel_case_types)]
125 #derive
126 struct #args { #(#fields),* }
127 },
128 pack: quote! { &#args { #(#names),* } },
129 unpack: quote! { #(a.#names),* },
130 }
131}
132
133fn output_type(sig: &syn::Signature) -> proc_macro2::TokenStream {
134 match &sig.output {
135 syn::ReturnType::Default => quote! { () },
136 syn::ReturnType::Type(_, ty) => quote! { #ty },
137 }
138}
139
140fn doc_of(attrs: &[syn::Attribute]) -> String {
144 let lines: Vec<String> = attrs
145 .iter()
146 .filter(|a| a.path().is_ident("doc"))
147 .filter_map(|a| match &a.meta {
148 syn::Meta::NameValue(nv) => match &nv.value {
149 syn::Expr::Lit(syn::ExprLit {
150 lit: syn::Lit::Str(s),
151 ..
152 }) => Some(s.value().trim().to_owned()),
153 _ => None,
154 },
155 _ => None,
156 })
157 .collect();
158 lines.join(" ")
159}
160
161fn implements_of(attr: TokenStream) -> syn::Result<Vec<String>> {
163 let mut out = Vec::new();
164 let parser = syn::meta::parser(|meta| {
165 if meta.path.is_ident("implements") {
166 out.push(meta.value()?.parse::<syn::LitStr>()?.value());
167 Ok(())
168 } else {
169 Err(meta.error("unknown `yog_export` argument (implements)"))
170 }
171 });
172 syn::parse::Parser::parse(parser, attr)?;
173 Ok(out)
174}
175
176#[proc_macro_attribute]
177pub fn yog_export(attr: TokenStream, item: TokenStream) -> TokenStream {
178 let implements = match implements_of(attr) {
179 Ok(i) => i,
180 Err(e) => return e.to_compile_error().into(),
181 };
182 let tokens = proc_macro2::TokenStream::from(item.clone());
183 if syn::parse2::<syn::ItemStruct>(tokens.clone()).is_ok() {
184 let derive = schema_derive();
185 return TokenStream::from(quote! { #derive #tokens });
186 }
187 if let Ok(e) = syn::parse2::<syn::ItemEnum>(tokens.clone()) {
188 let derive = schema_derive();
189 let tag = enum_tag(&e);
190 return TokenStream::from(quote! { #derive #tag #tokens });
191 }
192 let func = syn::parse_macro_input!(item as syn::ItemFn);
193 let name = &func.sig.ident;
194 let name_str = name.to_string();
195 let serve_name = format_ident!("__yog_serve_{}", name);
196 let ctor_name = format_ident!("__yog_ctor_register_{}", name);
197 let (registry, params) = call_params(&func.sig);
198 let Args {
199 ty: args,
200 def: args_def,
201 unpack,
202 ..
203 } = args_of(name, ¶ms);
204 let output = output_type(&func.sig);
205 let doc = doc_of(&func.attrs);
206 let (setup, registry_arg) = if registry {
208 (
209 quote! { let mut __yog_registry = unsafe { ::yog_api::Registry::from_raw(::yog_api::__current_api_ptr() as *const ::yog_api::YogApi) }; },
210 quote! { &mut __yog_registry, },
211 )
212 } else {
213 (quote! {}, quote! {})
214 };
215 TokenStream::from(quote! {
216 #func
217
218 #args_def
219
220 #[doc(hidden)]
221 unsafe extern "C" fn #serve_name(
222 _ud: *mut ::std::os::raw::c_void,
223 payload: ::yog_api::YogBytes,
224 reply: *const ::yog_api::YogReply,
225 ) -> ::yog_api::YogStatus {
226 unsafe {
227 ::yog_api::__yog_serve::<#args, #output>(payload, reply, |a| {
228 #setup
229 #name(#registry_arg #unpack)
230 })
231 }
232 }
233
234 #[::yog_api::ctor::ctor(unsafe, crate_path = ::yog_api::ctor)]
236 #[doc(hidden)]
237 fn #ctor_name() {
238 ::yog_api::__yog_export_registry().lock().push(
239 ::yog_api::YogExportEntry {
240 name: #name_str,
241 handler: #serve_name,
242 doc: #doc,
243 shapes: || (<#args as ::yog_api::yog_data::Schema>::shape(), <#output as ::yog_api::yog_data::Schema>::shape()),
244 implements: &[#(#implements),*],
245 }
246 );
247 }
248 })
249}
250
251#[proc_macro_attribute]
258pub fn bench(_attr: TokenStream, item: TokenStream) -> TokenStream {
259 let mut f = syn::parse_macro_input!(item as syn::ItemFn);
260 if matches!(f.vis, syn::Visibility::Inherited) {
261 f.vis = syn::parse_quote!(pub(crate));
262 }
263 TokenStream::from(quote! { #f })
264}
265
266#[proc_macro]
273pub fn import(input: TokenStream) -> TokenStream {
274 let import_block = syn::parse_macro_input!(input as ImportBlock);
275 let mod_name = &import_block.mod_name;
276 let mod_ident = format_ident!("{}", mod_name.replace('-', "_"));
277 let mut items = proc_macro2::TokenStream::new();
278 for item in &import_block.items {
279 match item {
280 ImportItem::Func(func) => items.extend(generate_import_fn(func, mod_name)),
281 ImportItem::Type(t) => {
282 let derive = schema_derive();
283 items.extend(quote! { #derive #t });
284 }
285 }
286 }
287 TokenStream::from(quote! { pub mod #mod_ident { use super::*; #items } })
288}
289
290struct ImportFn {
292 vis: syn::Visibility,
293 sig: syn::Signature,
294}
295
296fn generate_import_fn(func: &ImportFn, mod_name: &str) -> proc_macro2::TokenStream {
297 let name = &func.sig.ident;
298 let name_str = name.to_string();
299 let vis = &func.vis;
300 let (_, params) = call_params(&func.sig);
301 let Args {
302 ty: args,
303 def: args_def,
304 pack,
305 ..
306 } = args_of(name, ¶ms);
307 let output = output_type(&func.sig);
308 let names: Vec<&syn::Ident> = params.iter().map(|(n, _)| n).collect();
309 let types: Vec<&syn::Type> = params.iter().map(|(_, t)| t).collect();
310 quote! {
311 #args_def
312
313 #vis fn #name(#(#names: #types),*) -> ::std::option::Option<#output> {
315 ::yog_api::call_export::<#args, #output>(#mod_name, #name_str, #pack)
316 }
317 }
318}
319
320struct ImportBlock {
323 mod_name: String,
324 items: Vec<ImportItem>,
325}
326enum ImportItem {
327 Func(Box<ImportFn>),
328 Type(proc_macro2::TokenStream),
330}
331
332impl syn::parse::Parse for ImportBlock {
333 fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
334 let from_kw: syn::Ident = input.parse()?;
335 if from_kw != "from" {
336 return Err(syn::Error::new(from_kw.span(), "expected `from`"));
337 }
338 let lit: syn::LitStr = input.parse()?;
339 let mod_name = lit.value();
340 let content;
341 syn::braced!(content in input);
342 let mut items = Vec::new();
343 while !content.is_empty() {
344 let attrs = content.call(syn::Attribute::parse_outer)?;
345 let vis: syn::Visibility = content.parse()?;
346 if content.peek(syn::Token![fn]) {
347 let sig: syn::Signature = content.parse()?;
348 if content.peek(syn::token::Brace) {
349 let _body: syn::Block = content.parse()?;
350 } else {
351 content.parse::<syn::Token![;]>()?;
352 }
353 items.push(ImportItem::Func(Box::new(ImportFn { vis, sig })));
354 } else if content.peek(syn::Token![enum]) {
355 let mut e: syn::ItemEnum = content.parse()?;
356 e.attrs = attrs;
357 e.vis = vis;
358 items.push(ImportItem::Type(quote! { #e }));
359 } else {
360 let mut s: syn::ItemStruct = content.parse()?;
361 s.attrs = attrs;
362 s.vis = vis;
363 items.push(ImportItem::Type(quote! { #s }));
364 let _ = content.parse::<syn::Token![;]>();
365 }
366 }
367 Ok(ImportBlock { mod_name, items })
368 }
369}