use proc_macro::TokenStream;
use quote::quote;
use syn::{
Error, FnArg, ItemFn, LitStr, Meta, Pat, PatIdent, ReturnType, Token, Type, parse_macro_input,
punctuated::Punctuated,
};
#[proc_macro_attribute]
pub fn pd_host_function(attr: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attr with Punctuated::<Meta, Token![,]>::parse_terminated);
match expand_pd_host_function(args, parse_macro_input!(item as ItemFn)) {
Ok(tokens) => tokens.into(),
Err(err) => err.to_compile_error().into(),
}
}
fn expand_pd_host_function(
attr: Punctuated<Meta, Token![,]>,
mut item: ItemFn,
) -> Result<proc_macro2::TokenStream, Error> {
parse_name_arg(&attr)?;
let docs = doc_string(&item.attrs);
for input in &item.sig.inputs {
validate_param(input)?;
}
validate_return_type(&item.sig.output)?;
if is_abi_declaration_only(&item) {
return Ok(quote!(#item));
}
if docs.trim().is_empty() {
return Err(Error::new_spanned(
&item.sig.ident,
"#[pd_host_function] requires /// doc comments",
));
}
let (wrapper_name, impl_name) = wrapper_and_impl_names(&item.sig.ident);
if item.sig.ident != impl_name {
item.sig.ident = impl_name.clone();
}
let wrapper = generate_vm_wrapper(&item, &wrapper_name)?;
Ok(quote! {
#item
#wrapper
})
}
fn parse_name_arg(args: &Punctuated<Meta, Token![,]>) -> Result<LitStr, Error> {
let Some(Meta::NameValue(name_value)) = args.first() else {
return Err(Error::new(
proc_macro2::Span::call_site(),
"expected #[pd_host_function(name = \"...\")]",
));
};
if !name_value.path.is_ident("name") {
return Err(Error::new_spanned(
&name_value.path,
"expected #[pd_host_function(name = \"...\")]",
));
}
match &name_value.value {
syn::Expr::Lit(expr_lit) => {
if let syn::Lit::Str(value) = &expr_lit.lit {
Ok(value.clone())
} else {
Err(Error::new_spanned(
&expr_lit.lit,
"callable name must be a string literal",
))
}
}
other => Err(Error::new_spanned(
other,
"callable name must be a string literal",
)),
}
}
fn doc_string(attrs: &[syn::Attribute]) -> String {
attrs
.iter()
.filter_map(|attr| {
if !attr.path().is_ident("doc") {
return None;
}
match &attr.meta {
Meta::NameValue(name_value) => match &name_value.value {
syn::Expr::Lit(expr_lit) => match &expr_lit.lit {
syn::Lit::Str(value) => Some(value.value().trim().to_string()),
_ => None,
},
_ => None,
},
_ => None,
}
})
.filter(|line| !line.is_empty())
.collect::<Vec<_>>()
.join("\n")
}
fn validate_param(arg: &FnArg) -> Result<(), Error> {
let FnArg::Typed(pat_type) = arg else {
return Err(Error::new_spanned(arg, "methods are not supported"));
};
if is_vm_context_type(&pat_type.ty) {
return Ok(());
}
let Pat::Ident(PatIdent { .. }) = pat_type.pat.as_ref() else {
return Err(Error::new_spanned(
&pat_type.pat,
"callable parameters must use identifier patterns",
));
};
type_label(&pat_type.ty)?;
Ok(())
}
fn validate_return_type(output: &ReturnType) -> Result<(), Error> {
match output {
ReturnType::Default => Ok(()),
ReturnType::Type(_, ty) => {
type_label(ty)?;
Ok(())
}
}
}
fn is_abi_declaration_only(item: &ItemFn) -> bool {
let [stmt] = item.block.stmts.as_slice() else {
return false;
};
let syn::Stmt::Expr(expr, None) = stmt else {
return false;
};
let syn::Expr::Macro(expr_macro) = expr else {
return false;
};
expr_macro.mac.path.is_ident("unreachable")
}
fn generate_vm_wrapper(
item: &ItemFn,
wrapper_name: &syn::Ident,
) -> Result<proc_macro2::TokenStream, Error> {
let impl_name = &item.sig.ident;
let mut wrapper_params = Vec::<proc_macro2::TokenStream>::new();
let mut call_args = Vec::<proc_macro2::TokenStream>::new();
let mut imm_extract_stmts = Vec::<proc_macro2::TokenStream>::new();
let mut mut_extract_stmts = Vec::<proc_macro2::TokenStream>::new();
let mutable_wrapper_name = syn::Ident::new(&format!("{wrapper_name}_mut"), wrapper_name.span());
let has_vm = item.sig.inputs.iter().any(|input| match input {
FnArg::Typed(pat_type) => is_vm_context_type(&pat_type.ty),
FnArg::Receiver(_) => false,
});
if has_vm {
wrapper_params.push(quote!(vm: &mut super::super::Vm));
call_args.push(quote!(vm));
}
let imm_wrapper_params = {
let mut params = wrapper_params.clone();
params.push(quote!(args: &[super::super::Value]));
params
};
let mut_wrapper_params = {
let mut params = wrapper_params.clone();
params.push(quote!(args: &mut [super::super::Value]));
params
};
let mut arg_index = 0usize;
for input in &item.sig.inputs {
let FnArg::Typed(pat_type) = input else {
return Err(Error::new_spanned(input, "methods are not supported"));
};
if is_vm_context_type(&pat_type.ty) {
continue;
}
let Pat::Ident(PatIdent { ident, .. }) = pat_type.pat.as_ref() else {
return Err(Error::new_spanned(
&pat_type.pat,
"callable parameters must use identifier patterns",
));
};
let ty = &pat_type.ty;
let label = LitStr::new(
&format!("{} {}", wrapper_name, ident),
proc_macro2::Span::call_site(),
);
let index = syn::Index::from(arg_index);
imm_extract_stmts.push(quote! {
let #ident = super::borrow_arg::<#ty>(args, #index, #label)?;
});
let extractor = if uses_taken_extractor(ty) {
quote!(super::take_arg::<#ty>(args, #index, #label)?)
} else {
quote!(super::borrow_arg::<#ty>(&*args, #index, #label)?)
};
mut_extract_stmts.push(quote! {
let #ident = #extractor;
});
call_args.push(quote!(#ident));
arg_index += 1;
}
let wrapper_output = wrapper_output_type(&item.sig.output)?;
let call_expr = if return_is_vm_result(&item.sig.output) {
quote!(#impl_name(#(#call_args),*))
} else {
quote!(Ok(#impl_name(#(#call_args),*)))
};
Ok(quote! {
#[allow(dead_code)]
pub(super) fn #wrapper_name(#(#imm_wrapper_params),*) -> #wrapper_output {
#(#imm_extract_stmts)*
#call_expr
}
#[allow(dead_code)]
pub(super) fn #mutable_wrapper_name(#(#mut_wrapper_params),*) -> #wrapper_output {
#(#mut_extract_stmts)*
#call_expr
}
})
}
fn wrapper_and_impl_names(name: &syn::Ident) -> (syn::Ident, syn::Ident) {
let original = name.to_string();
match original.strip_suffix("_impl") {
Some(prefix) => (
syn::Ident::new(prefix, name.span()),
syn::Ident::new(&original, name.span()),
),
None => (
syn::Ident::new(&original, name.span()),
syn::Ident::new(&format!("{original}_impl"), name.span()),
),
}
}
fn wrapper_output_type(output: &ReturnType) -> Result<proc_macro2::TokenStream, Error> {
if let Some(inner) = vm_result_inner_type(output)? {
return Ok(quote!(super::super::VmResult<#inner>));
}
match output {
ReturnType::Default => Ok(quote!(super::super::VmResult<()>)),
ReturnType::Type(_, ty) => Ok(quote!(super::super::VmResult<#ty>)),
}
}
fn vm_result_inner_type(output: &ReturnType) -> Result<Option<Type>, Error> {
let ReturnType::Type(_, ty) = output else {
return Ok(None);
};
unwrap_vm_result_type(ty)
}
fn unwrap_vm_result_type(ty: &Type) -> Result<Option<Type>, Error> {
match ty {
Type::Group(group) => unwrap_vm_result_type(&group.elem),
Type::Paren(paren) => unwrap_vm_result_type(&paren.elem),
Type::Reference(reference) => unwrap_vm_result_type(&reference.elem),
Type::Path(path) => {
let Some(segment) = path.path.segments.last() else {
return Ok(None);
};
if !matches!(
segment.ident.to_string().as_str(),
"VmResult" | "HostResult"
) {
return Ok(None);
}
let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
return Err(Error::new_spanned(
&segment.arguments,
format!("{}<T> requires one generic argument", segment.ident),
));
};
let Some(syn::GenericArgument::Type(inner)) = args.args.first() else {
return Err(Error::new_spanned(
args,
format!("{}<T> requires one type argument", segment.ident),
));
};
Ok(Some(inner.clone()))
}
_ => Ok(None),
}
}
fn return_is_vm_result(output: &ReturnType) -> bool {
vm_result_inner_type(output)
.expect("pd_host_function return type should already be validated")
.is_some()
}
fn type_label(ty: &Type) -> Result<String, Error> {
match ty {
Type::Group(group) => type_label(&group.elem),
Type::Paren(paren) => type_label(&paren.elem),
Type::Reference(reference) => type_label(&reference.elem),
Type::Slice(slice) => match slice.elem.as_ref() {
Type::Path(path) => {
let Some(segment) = path.path.segments.last() else {
return Err(Error::new_spanned(slice, "unsupported callable type"));
};
if segment.ident == "u8" {
Ok("bytes".to_string())
} else {
Err(Error::new_spanned(slice, "unsupported callable type"))
}
}
_ => Err(Error::new_spanned(slice, "unsupported callable type")),
},
Type::Tuple(tuple) if tuple.elems.is_empty() => Ok("null".to_string()),
Type::Path(path) => {
let Some(segment) = path.path.segments.last() else {
return Err(Error::new_spanned(path, "unsupported callable type"));
};
let ident = segment.ident.to_string();
match ident.as_str() {
"i8" | "i16" | "i32" | "i64" | "i128" | "isize" | "u8" | "u16" | "u32" | "u64"
| "u128" | "usize" => Ok("int".to_string()),
"f32" | "f64" => Ok("float".to_string()),
"bool" => Ok("bool".to_string()),
"String" | "str" | "VmStringRef" => Ok("string".to_string()),
"Bytes" | "VmBytes" | "VmBytesRef" | "VmBytesHandle" => Ok("bytes".to_string()),
"Any" | "AnyValue" | "Value" | "VmValueRef" | "VmValueOwned" => {
Ok("any".to_string())
}
"Array" | "VmArray" | "VmArrayRef" | "VmArrayHandle" => Ok("array".to_string()),
"Map" | "VmMap" | "VmMapRef" | "VmMapHandle" => Ok("map".to_string()),
"Number" | "NumberValue" => Ok("number".to_string()),
"Unknown" | "UnknownValue" => Ok("unknown".to_string()),
"CallOutcome" => Ok("unknown".to_string()),
"Option" => {
let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
return Err(Error::new_spanned(
&segment.arguments,
"Option<T> requires one generic argument",
));
};
let Some(syn::GenericArgument::Type(inner)) = args.args.first() else {
return Err(Error::new_spanned(
args,
"Option<T> requires one type argument",
));
};
let inner_label = type_label(inner)?;
Ok(format!("{inner_label} | null"))
}
"VmResult" | "BuiltinResult" | "HostResult" => {
let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
return Err(Error::new_spanned(
&segment.arguments,
format!("{ident}<T> requires one generic argument"),
));
};
let Some(syn::GenericArgument::Type(inner)) = args.args.first() else {
return Err(Error::new_spanned(
args,
format!("{ident}<T> requires one type argument"),
));
};
type_label(inner)
}
"Vec" => type_label_for_vec(segment),
_ => Err(Error::new_spanned(
path,
format!("unsupported callable type '{ident}'"),
)),
}
}
_ => Err(Error::new_spanned(ty, "unsupported callable type")),
}
}
fn type_label_for_vec(segment: &syn::PathSegment) -> Result<String, Error> {
let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
return Err(Error::new_spanned(
&segment.arguments,
"Vec<T> requires one generic argument",
));
};
let Some(syn::GenericArgument::Type(inner)) = args.args.first() else {
return Err(Error::new_spanned(
args,
"Vec<T> requires one type argument",
));
};
match inner {
Type::Tuple(tuple) if tuple.elems.len() == 2 => {
let lhs = tuple
.elems
.first()
.expect("tuple should contain first element");
let rhs = tuple
.elems
.last()
.expect("tuple should contain second element");
if is_value_type(lhs) && is_value_type(rhs) {
Ok("map".to_string())
} else {
Err(Error::new_spanned(
inner,
"unsupported Vec tuple type in callable metadata",
))
}
}
_ if is_value_type(inner) => Ok("array".to_string()),
_ => {
let inner_label = type_label(inner)?;
Err(Error::new_spanned(
inner,
format!("unsupported Vec return type '{inner_label}'"),
))
}
}
}
fn is_value_type(ty: &Type) -> bool {
match ty {
Type::Group(group) => is_value_type(&group.elem),
Type::Paren(paren) => is_value_type(&paren.elem),
Type::Reference(reference) => is_value_type(&reference.elem),
Type::Path(path) => path
.path
.segments
.last()
.is_some_and(|segment| segment.ident == "Value"),
_ => false,
}
}
fn is_vm_context_type(ty: &Type) -> bool {
match ty {
Type::Group(group) => is_vm_context_type(&group.elem),
Type::Paren(paren) => is_vm_context_type(&paren.elem),
Type::Reference(reference) => is_vm_context_type(&reference.elem),
Type::Path(path) => path
.path
.segments
.last()
.is_some_and(|segment| segment.ident == "Vm"),
_ => false,
}
}
fn uses_taken_extractor(ty: &Type) -> bool {
match ty {
Type::Group(group) => uses_taken_extractor(&group.elem),
Type::Paren(paren) => uses_taken_extractor(&paren.elem),
Type::Reference(_) => false,
Type::Path(path) => path.path.segments.last().is_some_and(|segment| {
matches!(
segment.ident.to_string().as_str(),
"Value"
| "AnyValue"
| "UnknownValue"
| "VmArray"
| "VmBytes"
| "VmMap"
| "VmArrayHandle"
| "VmBytesHandle"
| "VmMapHandle"
| "VmValueOwned"
)
}),
_ => false,
}
}