From d842d69be25954ff4d15f17b8bb1c5a2afcc2d89 Mon Sep 17 00:00:00 2001 From: Gino Valente Date: Mon, 9 Dec 2024 22:09:03 -0700 Subject: [PATCH] Add reflect_fn proc macro --- .../bevy_reflect/derive/src/impls/func/mod.rs | 2 + .../derive/src/impls/func/reflect_fn.rs | 167 ++++++++++++++++++ crates/bevy_reflect/derive/src/impls/mod.rs | 2 +- crates/bevy_reflect/derive/src/lib.rs | 6 + crates/bevy_reflect/src/func/mod.rs | 76 ++++++++ 5 files changed, 252 insertions(+), 1 deletion(-) create mode 100644 crates/bevy_reflect/derive/src/impls/func/reflect_fn.rs diff --git a/crates/bevy_reflect/derive/src/impls/func/mod.rs b/crates/bevy_reflect/derive/src/impls/func/mod.rs index b092366146..a3e6e9aa64 100644 --- a/crates/bevy_reflect/derive/src/impls/func/mod.rs +++ b/crates/bevy_reflect/derive/src/impls/func/mod.rs @@ -1,6 +1,8 @@ pub(crate) use function_impls::impl_function_traits; +pub(crate) use reflect_fn::reflect_fn; mod from_arg; mod function_impls; mod get_ownership; mod into_return; +mod reflect_fn; diff --git a/crates/bevy_reflect/derive/src/impls/func/reflect_fn.rs b/crates/bevy_reflect/derive/src/impls/func/reflect_fn.rs new file mode 100644 index 0000000000..d3f1f31284 --- /dev/null +++ b/crates/bevy_reflect/derive/src/impls/func/reflect_fn.rs @@ -0,0 +1,167 @@ +use bevy_macro_utils::fq_std::FQResult; +use proc_macro2::Ident; +use quote::{format_ident, quote}; +use syn::parse::{Parse, ParseStream}; +use syn::punctuated::Punctuated; +use syn::{parenthesized, parse_macro_input, token, Block, ExprBlock, Type}; +use syn::{LitStr, ReturnType, Token}; + +struct Arg { + mutability: Option, + name: Ident, + _colon: Token![:], + ty: Type, +} + +impl Parse for Arg { + fn parse(input: ParseStream) -> syn::Result { + let mutability = if input.peek(Token![mut]) { + Some(input.parse()?) + } else { + None + }; + + Ok(Self { + mutability, + name: input.parse()?, + _colon: input.parse()?, + ty: input.parse()?, + }) + } +} + +enum FnName { + #[expect( + dead_code, + reason = "for documenting via the type system what `Anon` expects to parse" + )] + Anon(Token![_]), + Lit(LitStr), + Name(Ident), + Expr(ExprBlock), +} + +impl Parse for FnName { + fn parse(input: ParseStream) -> syn::Result { + let lookahead = input.lookahead1(); + if lookahead.peek(Token![_]) { + Ok(Self::Anon(input.parse()?)) + } else if lookahead.peek(LitStr) { + Ok(Self::Lit(input.parse()?)) + } else if lookahead.peek(syn::Ident) { + Ok(Self::Name(input.parse()?)) + } else if lookahead.peek(token::Brace) { + Ok(Self::Expr(input.parse()?)) + } else { + Err(lookahead.error()) + } + } +} + +struct ReflectFn { + mutability: Option, + movability: Option, + _fn_token: Token![fn], + name: FnName, + _parens: token::Paren, + args: Punctuated, + return_type: ReturnType, + body: Block, +} + +impl Parse for ReflectFn { + fn parse(input: ParseStream) -> syn::Result { + let mutability = if input.peek(Token![mut]) { + Some(input.parse()?) + } else { + None + }; + let movability = if input.peek(Token![move]) { + Some(input.parse()?) + } else { + None + }; + + let content; + Ok(Self { + mutability, + movability, + _fn_token: input.parse()?, + name: input.parse()?, + _parens: parenthesized!(content in input), + args: content.parse_terminated(Arg::parse, Token![,])?, + return_type: input.parse()?, + body: input.parse()?, + }) + } +} + +pub(crate) fn reflect_fn(input: proc_macro::TokenStream) -> proc_macro::TokenStream { + let input = parse_macro_input!(input as ReflectFn); + + let bevy_reflect_path = crate::meta::get_bevy_reflect_path(); + + let arg_list = format_ident!("args"); + + let dynamic_function = if input.mutability.is_some() { + quote! { #bevy_reflect_path::func::DynamicFunctionMut } + } else { + quote! { #bevy_reflect_path::func::DynamicFunction } + }; + + let movability = &input.movability; + + let extract_args = input.args.iter().map(|arg| { + let mutability = &arg.mutability; + let name = &arg.name; + let ty = &arg.ty; + + quote! { let #mutability #name = #arg_list.take::<#ty>()?; } + }); + + let body = &input.body; + + let info = match &input.name { + FnName::Anon(_) => quote! { #bevy_reflect_path::func::SignatureInfo::anonymous() }, + FnName::Lit(name) => { + quote! { #bevy_reflect_path::func::SignatureInfo::named(#name) } + } + FnName::Name(name) => { + let name = name.to_string(); + quote! { #bevy_reflect_path::func::SignatureInfo::named(#name) } + } + FnName::Expr(expr) => { + quote! { #bevy_reflect_path::func::SignatureInfo::named(#expr) } + } + }; + + let arg_info = input.args.iter().enumerate().map(|(index, arg)| { + let name = &arg.name; + let ty = &arg.ty; + + quote! { + #bevy_reflect_path::func::args::ArgInfo::new::<#ty>(#index).with_name(stringify!(#name)) + } + }); + + let return_ty = match input.return_type { + ReturnType::Default => quote! { () }, + ReturnType::Type(_, ty) => quote! { #ty }, + }; + + proc_macro::TokenStream::from(quote! {{ + #dynamic_function::new( + #[allow(unused_mut)] + #movability |mut #arg_list| { + #(#extract_args)* + #FQResult::Ok(#bevy_reflect_path::func::IntoReturn::into_return(#body)) + }, + #[allow(unused_braces)] + #bevy_reflect_path::func::FunctionInfo::new( + #info + .with_args(::alloc::vec![#(#arg_info),*]) + .with_return::<#return_ty>() + ), + ) + }}) +} diff --git a/crates/bevy_reflect/derive/src/impls/mod.rs b/crates/bevy_reflect/derive/src/impls/mod.rs index 6477c4041e..f1ea786858 100644 --- a/crates/bevy_reflect/derive/src/impls/mod.rs +++ b/crates/bevy_reflect/derive/src/impls/mod.rs @@ -2,7 +2,7 @@ mod assertions; mod common; mod enums; #[cfg(feature = "functions")] -mod func; +pub(crate) mod func; mod opaque; mod structs; mod tuple_structs; diff --git a/crates/bevy_reflect/derive/src/lib.rs b/crates/bevy_reflect/derive/src/lib.rs index 276371427b..9b02bc3f16 100644 --- a/crates/bevy_reflect/derive/src/lib.rs +++ b/crates/bevy_reflect/derive/src/lib.rs @@ -826,3 +826,9 @@ pub fn impl_type_path(input: TokenStream) -> TokenStream { }; }) } + +#[cfg(feature = "functions")] +#[proc_macro] +pub fn reflect_fn(input: TokenStream) -> TokenStream { + impls::func::reflect_fn(input) +} diff --git a/crates/bevy_reflect/src/func/mod.rs b/crates/bevy_reflect/src/func/mod.rs index f990e135f0..e27a435e1b 100644 --- a/crates/bevy_reflect/src/func/mod.rs +++ b/crates/bevy_reflect/src/func/mod.rs @@ -189,6 +189,7 @@ pub mod signature; #[cfg(test)] mod tests { + use crate as bevy_reflect; use alloc::borrow::Cow; use super::*; @@ -197,6 +198,7 @@ mod tests { func::args::{ArgError, ArgList, Ownership}, TypePath, }; + use bevy_reflect_derive::reflect_fn; #[test] fn should_error_on_missing_args() { @@ -263,4 +265,78 @@ mod tests { }) ); } + + #[test] + fn should_create_dynamic_function_with_macro() { + let func = reflect_fn!( + fn add(a: i32, b: i32) -> i32 { + a + b + } + ); + + let info = func.info(); + + assert_eq!(info.name().unwrap(), "add"); + assert_eq!(info.base().arg_count(), 2); + assert_eq!(info.base().args()[0].name(), Some("a")); + assert!(info.base().args()[0].is::()); + assert_eq!(info.base().args()[1].name(), Some("b")); + assert!(info.base().args()[1].is::()); + assert!(info.base().return_info().is::()); + + let args = ArgList::new().push_owned(25_i32).push_owned(75_i32); + let result = func.call(args).unwrap().unwrap_owned(); + assert_eq!(result.try_downcast_ref::(), Some(&100)); + } + + #[test] + fn should_create_zero_arg_dynamic_function_with_macro() { + let func = reflect_fn!( + fn shout() { + println!("HEY!!!"); + } + ); + + let info = func.info(); + + assert_eq!(info.name().unwrap(), "shout"); + assert_eq!(info.base().arg_count(), 0); + assert!(info.base().return_info().is::<()>()); + } + + #[test] + fn should_create_anonymous_dynamic_function_with_macro() { + let func = reflect_fn!( + fn _(a: i32, b: i32) -> i32 { + a + b + } + ); + + let info = func.info(); + assert_eq!(info.name(), None); + } + + #[test] + fn should_create_string_named_dynamic_function_with_macro() { + let func = reflect_fn!( + fn "some cool function"(a: i32, b: i32) -> i32 { + a + b + } + ); + + let info = func.info(); + assert_eq!(info.name().unwrap(), "some cool function"); + } + + #[test] + fn should_create_expression_named_dynamic_function_with_macro() { + let func = reflect_fn!( + fn {concat!("some", " cool ", "function")}(a: i32, b: i32) -> i32 { + a + b + } + ); + + let info = func.info(); + assert_eq!(info.name().unwrap(), "some cool function"); + } }