pd_host_function/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{
4    Error, FnArg, ItemFn, LitStr, Meta, Pat, PatIdent, ReturnType, Token, Type, parse_macro_input,
5    punctuated::Punctuated,
6};
7
8#[proc_macro_attribute]
9pub fn pd_host_function(attr: TokenStream, item: TokenStream) -> TokenStream {
10    let args = parse_macro_input!(attr with Punctuated::<Meta, Token![,]>::parse_terminated);
11    match expand_pd_host_function(args, parse_macro_input!(item as ItemFn)) {
12        Ok(tokens) => tokens.into(),
13        Err(err) => err.to_compile_error().into(),
14    }
15}
16
17fn expand_pd_host_function(
18    attr: Punctuated<Meta, Token![,]>,
19    mut item: ItemFn,
20) -> Result<proc_macro2::TokenStream, Error> {
21    parse_name_arg(&attr)?;
22    let docs = doc_string(&item.attrs);
23    for input in &item.sig.inputs {
24        validate_param(input)?;
25    }
26    validate_return_type(&item.sig.output)?;
27
28    if is_abi_declaration_only(&item) {
29        return Ok(quote!(#item));
30    }
31    if docs.trim().is_empty() {
32        return Err(Error::new_spanned(
33            &item.sig.ident,
34            "#[pd_host_function] requires /// doc comments",
35        ));
36    }
37
38    let (wrapper_name, impl_name) = wrapper_and_impl_names(&item.sig.ident);
39    if item.sig.ident != impl_name {
40        item.sig.ident = impl_name.clone();
41    }
42    let wrapper = generate_vm_wrapper(&item, &wrapper_name)?;
43    Ok(quote! {
44        #item
45        #wrapper
46    })
47}
48
49fn parse_name_arg(args: &Punctuated<Meta, Token![,]>) -> Result<LitStr, Error> {
50    let Some(Meta::NameValue(name_value)) = args.first() else {
51        return Err(Error::new(
52            proc_macro2::Span::call_site(),
53            "expected #[pd_host_function(name = \"...\")]",
54        ));
55    };
56    if !name_value.path.is_ident("name") {
57        return Err(Error::new_spanned(
58            &name_value.path,
59            "expected #[pd_host_function(name = \"...\")]",
60        ));
61    }
62    match &name_value.value {
63        syn::Expr::Lit(expr_lit) => {
64            if let syn::Lit::Str(value) = &expr_lit.lit {
65                Ok(value.clone())
66            } else {
67                Err(Error::new_spanned(
68                    &expr_lit.lit,
69                    "callable name must be a string literal",
70                ))
71            }
72        }
73        other => Err(Error::new_spanned(
74            other,
75            "callable name must be a string literal",
76        )),
77    }
78}
79
80fn doc_string(attrs: &[syn::Attribute]) -> String {
81    attrs
82        .iter()
83        .filter_map(|attr| {
84            if !attr.path().is_ident("doc") {
85                return None;
86            }
87            match &attr.meta {
88                Meta::NameValue(name_value) => match &name_value.value {
89                    syn::Expr::Lit(expr_lit) => match &expr_lit.lit {
90                        syn::Lit::Str(value) => Some(value.value().trim().to_string()),
91                        _ => None,
92                    },
93                    _ => None,
94                },
95                _ => None,
96            }
97        })
98        .filter(|line| !line.is_empty())
99        .collect::<Vec<_>>()
100        .join("\n")
101}
102
103fn validate_param(arg: &FnArg) -> Result<(), Error> {
104    let FnArg::Typed(pat_type) = arg else {
105        return Err(Error::new_spanned(arg, "methods are not supported"));
106    };
107    if is_vm_context_type(&pat_type.ty) {
108        return Ok(());
109    }
110    let Pat::Ident(PatIdent { .. }) = pat_type.pat.as_ref() else {
111        return Err(Error::new_spanned(
112            &pat_type.pat,
113            "callable parameters must use identifier patterns",
114        ));
115    };
116    type_label(&pat_type.ty)?;
117    Ok(())
118}
119
120fn validate_return_type(output: &ReturnType) -> Result<(), Error> {
121    match output {
122        ReturnType::Default => Ok(()),
123        ReturnType::Type(_, ty) => {
124            type_label(ty)?;
125            Ok(())
126        }
127    }
128}
129
130fn is_abi_declaration_only(item: &ItemFn) -> bool {
131    let [stmt] = item.block.stmts.as_slice() else {
132        return false;
133    };
134    let syn::Stmt::Expr(expr, None) = stmt else {
135        return false;
136    };
137    let syn::Expr::Macro(expr_macro) = expr else {
138        return false;
139    };
140    expr_macro.mac.path.is_ident("unreachable")
141}
142
143fn generate_vm_wrapper(
144    item: &ItemFn,
145    wrapper_name: &syn::Ident,
146) -> Result<proc_macro2::TokenStream, Error> {
147    let impl_name = &item.sig.ident;
148    let mut wrapper_params = Vec::<proc_macro2::TokenStream>::new();
149    let mut call_args = Vec::<proc_macro2::TokenStream>::new();
150    let mut imm_extract_stmts = Vec::<proc_macro2::TokenStream>::new();
151    let mut mut_extract_stmts = Vec::<proc_macro2::TokenStream>::new();
152    let mutable_wrapper_name = syn::Ident::new(&format!("{wrapper_name}_mut"), wrapper_name.span());
153    let has_vm = item.sig.inputs.iter().any(|input| match input {
154        FnArg::Typed(pat_type) => is_vm_context_type(&pat_type.ty),
155        FnArg::Receiver(_) => false,
156    });
157    if has_vm {
158        wrapper_params.push(quote!(vm: &mut super::super::Vm));
159        call_args.push(quote!(vm));
160    }
161    let imm_wrapper_params = {
162        let mut params = wrapper_params.clone();
163        params.push(quote!(args: &[super::super::Value]));
164        params
165    };
166    let mut_wrapper_params = {
167        let mut params = wrapper_params.clone();
168        params.push(quote!(args: &mut [super::super::Value]));
169        params
170    };
171
172    let mut arg_index = 0usize;
173    for input in &item.sig.inputs {
174        let FnArg::Typed(pat_type) = input else {
175            return Err(Error::new_spanned(input, "methods are not supported"));
176        };
177        if is_vm_context_type(&pat_type.ty) {
178            continue;
179        }
180        let Pat::Ident(PatIdent { ident, .. }) = pat_type.pat.as_ref() else {
181            return Err(Error::new_spanned(
182                &pat_type.pat,
183                "callable parameters must use identifier patterns",
184            ));
185        };
186        let ty = &pat_type.ty;
187        let label = LitStr::new(
188            &format!("{} {}", wrapper_name, ident),
189            proc_macro2::Span::call_site(),
190        );
191        let index = syn::Index::from(arg_index);
192        imm_extract_stmts.push(quote! {
193            let #ident = super::borrow_arg::<#ty>(args, #index, #label)?;
194        });
195        let extractor = if uses_taken_extractor(ty) {
196            quote!(super::take_arg::<#ty>(args, #index, #label)?)
197        } else {
198            quote!(super::borrow_arg::<#ty>(&*args, #index, #label)?)
199        };
200        mut_extract_stmts.push(quote! {
201            let #ident = #extractor;
202        });
203        call_args.push(quote!(#ident));
204        arg_index += 1;
205    }
206
207    let wrapper_output = wrapper_output_type(&item.sig.output)?;
208    let call_expr = if return_is_vm_result(&item.sig.output) {
209        quote!(#impl_name(#(#call_args),*))
210    } else {
211        quote!(Ok(#impl_name(#(#call_args),*)))
212    };
213
214    Ok(quote! {
215        #[allow(dead_code)]
216        pub(super) fn #wrapper_name(#(#imm_wrapper_params),*) -> #wrapper_output {
217            #(#imm_extract_stmts)*
218            #call_expr
219        }
220
221        #[allow(dead_code)]
222        pub(super) fn #mutable_wrapper_name(#(#mut_wrapper_params),*) -> #wrapper_output {
223            #(#mut_extract_stmts)*
224            #call_expr
225        }
226    })
227}
228
229fn wrapper_and_impl_names(name: &syn::Ident) -> (syn::Ident, syn::Ident) {
230    let original = name.to_string();
231    match original.strip_suffix("_impl") {
232        Some(prefix) => (
233            syn::Ident::new(prefix, name.span()),
234            syn::Ident::new(&original, name.span()),
235        ),
236        None => (
237            syn::Ident::new(&original, name.span()),
238            syn::Ident::new(&format!("{original}_impl"), name.span()),
239        ),
240    }
241}
242
243fn wrapper_output_type(output: &ReturnType) -> Result<proc_macro2::TokenStream, Error> {
244    if let Some(inner) = vm_result_inner_type(output)? {
245        return Ok(quote!(super::super::VmResult<#inner>));
246    }
247
248    match output {
249        ReturnType::Default => Ok(quote!(super::super::VmResult<()>)),
250        ReturnType::Type(_, ty) => Ok(quote!(super::super::VmResult<#ty>)),
251    }
252}
253
254fn vm_result_inner_type(output: &ReturnType) -> Result<Option<Type>, Error> {
255    let ReturnType::Type(_, ty) = output else {
256        return Ok(None);
257    };
258    unwrap_vm_result_type(ty)
259}
260
261fn unwrap_vm_result_type(ty: &Type) -> Result<Option<Type>, Error> {
262    match ty {
263        Type::Group(group) => unwrap_vm_result_type(&group.elem),
264        Type::Paren(paren) => unwrap_vm_result_type(&paren.elem),
265        Type::Reference(reference) => unwrap_vm_result_type(&reference.elem),
266        Type::Path(path) => {
267            let Some(segment) = path.path.segments.last() else {
268                return Ok(None);
269            };
270            if !matches!(
271                segment.ident.to_string().as_str(),
272                "VmResult" | "HostResult"
273            ) {
274                return Ok(None);
275            }
276            let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
277                return Err(Error::new_spanned(
278                    &segment.arguments,
279                    format!("{}<T> requires one generic argument", segment.ident),
280                ));
281            };
282            let Some(syn::GenericArgument::Type(inner)) = args.args.first() else {
283                return Err(Error::new_spanned(
284                    args,
285                    format!("{}<T> requires one type argument", segment.ident),
286                ));
287            };
288            Ok(Some(inner.clone()))
289        }
290        _ => Ok(None),
291    }
292}
293
294fn return_is_vm_result(output: &ReturnType) -> bool {
295    vm_result_inner_type(output)
296        .expect("pd_host_function return type should already be validated")
297        .is_some()
298}
299
300fn type_label(ty: &Type) -> Result<String, Error> {
301    match ty {
302        Type::Group(group) => type_label(&group.elem),
303        Type::Paren(paren) => type_label(&paren.elem),
304        Type::Reference(reference) => type_label(&reference.elem),
305        Type::Slice(slice) => match slice.elem.as_ref() {
306            Type::Path(path) => {
307                let Some(segment) = path.path.segments.last() else {
308                    return Err(Error::new_spanned(slice, "unsupported callable type"));
309                };
310                if segment.ident == "u8" {
311                    Ok("bytes".to_string())
312                } else {
313                    Err(Error::new_spanned(slice, "unsupported callable type"))
314                }
315            }
316            _ => Err(Error::new_spanned(slice, "unsupported callable type")),
317        },
318        Type::Tuple(tuple) if tuple.elems.is_empty() => Ok("null".to_string()),
319        Type::Path(path) => {
320            let Some(segment) = path.path.segments.last() else {
321                return Err(Error::new_spanned(path, "unsupported callable type"));
322            };
323            let ident = segment.ident.to_string();
324            match ident.as_str() {
325                "i8" | "i16" | "i32" | "i64" | "i128" | "isize" | "u8" | "u16" | "u32" | "u64"
326                | "u128" | "usize" => Ok("int".to_string()),
327                "f32" | "f64" => Ok("float".to_string()),
328                "bool" => Ok("bool".to_string()),
329                "String" | "str" | "VmStringRef" => Ok("string".to_string()),
330                "Bytes" | "VmBytes" | "VmBytesRef" | "VmBytesHandle" => Ok("bytes".to_string()),
331                "Any" | "AnyValue" | "Value" | "VmValueRef" | "VmValueOwned" => {
332                    Ok("any".to_string())
333                }
334                "Array" | "VmArray" | "VmArrayRef" | "VmArrayHandle" => Ok("array".to_string()),
335                "Map" | "VmMap" | "VmMapRef" | "VmMapHandle" => Ok("map".to_string()),
336                "Number" | "NumberValue" => Ok("number".to_string()),
337                "Unknown" | "UnknownValue" => Ok("unknown".to_string()),
338                "CallOutcome" => Ok("unknown".to_string()),
339                "Option" => {
340                    let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
341                        return Err(Error::new_spanned(
342                            &segment.arguments,
343                            "Option<T> requires one generic argument",
344                        ));
345                    };
346                    let Some(syn::GenericArgument::Type(inner)) = args.args.first() else {
347                        return Err(Error::new_spanned(
348                            args,
349                            "Option<T> requires one type argument",
350                        ));
351                    };
352                    let inner_label = type_label(inner)?;
353                    Ok(format!("{inner_label} | null"))
354                }
355                "VmResult" | "BuiltinResult" | "HostResult" => {
356                    let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
357                        return Err(Error::new_spanned(
358                            &segment.arguments,
359                            format!("{ident}<T> requires one generic argument"),
360                        ));
361                    };
362                    let Some(syn::GenericArgument::Type(inner)) = args.args.first() else {
363                        return Err(Error::new_spanned(
364                            args,
365                            format!("{ident}<T> requires one type argument"),
366                        ));
367                    };
368                    type_label(inner)
369                }
370                "Vec" => type_label_for_vec(segment),
371                _ => Err(Error::new_spanned(
372                    path,
373                    format!("unsupported callable type '{ident}'"),
374                )),
375            }
376        }
377        _ => Err(Error::new_spanned(ty, "unsupported callable type")),
378    }
379}
380
381fn type_label_for_vec(segment: &syn::PathSegment) -> Result<String, Error> {
382    let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
383        return Err(Error::new_spanned(
384            &segment.arguments,
385            "Vec<T> requires one generic argument",
386        ));
387    };
388    let Some(syn::GenericArgument::Type(inner)) = args.args.first() else {
389        return Err(Error::new_spanned(
390            args,
391            "Vec<T> requires one type argument",
392        ));
393    };
394    match inner {
395        Type::Tuple(tuple) if tuple.elems.len() == 2 => {
396            let lhs = tuple
397                .elems
398                .first()
399                .expect("tuple should contain first element");
400            let rhs = tuple
401                .elems
402                .last()
403                .expect("tuple should contain second element");
404            if is_value_type(lhs) && is_value_type(rhs) {
405                Ok("map".to_string())
406            } else {
407                Err(Error::new_spanned(
408                    inner,
409                    "unsupported Vec tuple type in callable metadata",
410                ))
411            }
412        }
413        _ if is_value_type(inner) => Ok("array".to_string()),
414        _ => {
415            let inner_label = type_label(inner)?;
416            Err(Error::new_spanned(
417                inner,
418                format!("unsupported Vec return type '{inner_label}'"),
419            ))
420        }
421    }
422}
423
424fn is_value_type(ty: &Type) -> bool {
425    match ty {
426        Type::Group(group) => is_value_type(&group.elem),
427        Type::Paren(paren) => is_value_type(&paren.elem),
428        Type::Reference(reference) => is_value_type(&reference.elem),
429        Type::Path(path) => path
430            .path
431            .segments
432            .last()
433            .is_some_and(|segment| segment.ident == "Value"),
434        _ => false,
435    }
436}
437
438fn is_vm_context_type(ty: &Type) -> bool {
439    match ty {
440        Type::Group(group) => is_vm_context_type(&group.elem),
441        Type::Paren(paren) => is_vm_context_type(&paren.elem),
442        Type::Reference(reference) => is_vm_context_type(&reference.elem),
443        Type::Path(path) => path
444            .path
445            .segments
446            .last()
447            .is_some_and(|segment| segment.ident == "Vm"),
448        _ => false,
449    }
450}
451
452fn uses_taken_extractor(ty: &Type) -> bool {
453    match ty {
454        Type::Group(group) => uses_taken_extractor(&group.elem),
455        Type::Paren(paren) => uses_taken_extractor(&paren.elem),
456        Type::Reference(_) => false,
457        Type::Path(path) => path.path.segments.last().is_some_and(|segment| {
458            matches!(
459                segment.ident.to_string().as_str(),
460                "Value"
461                    | "AnyValue"
462                    | "UnknownValue"
463                    | "VmArray"
464                    | "VmBytes"
465                    | "VmMap"
466                    | "VmArrayHandle"
467                    | "VmBytesHandle"
468                    | "VmMapHandle"
469                    | "VmValueOwned"
470            )
471        }),
472        _ => false,
473    }
474}
pd_host_function/lib.rs

pd_host_function/
lib.rs
