feat(engine-macros): improve flexibility of Reflection macro on generic typesHEADmaster

1 files changed, 272 insertions, 205 deletions

diff --git a/engine-macros/src/reflection/enum_impl.rs b/engine-macros/src/reflection/enum_impl.rs
index 5dbb176..0ce2562 100644
--- a/engine-macros/src/reflection/enum_impl.rs
+++ b/engine-macros/src/reflection/enum_impl.rs
@@ -1,37 +1,20 @@
 use quote::{format_ident, quote, ToTokens};
+use syn::ItemEnum;
 
 use crate::reflection::default_value::gen_get_default_value_fn;
 use crate::reflection::field::{generate as generate_field, ReflectionFieldGenOptions};
-use crate::util::{find_engine_crate_path, syn_path, SynPathExt};
+use crate::reflection::options_attr::OptionsAttr;
+use crate::util::find_engine_crate_path;
 
-const PRIMITIVE_REPRS: &[&'static str] = &[
-    "u8", "u16", "u32", "u64", "u128", "usize", "i8", "i16", "i32", "i64", "i128",
-    "isize",
-];
-
-pub fn generate(input: syn::ItemEnum) -> proc_macro2::TokenStream
+pub fn generate(input: syn::ItemEnum, options: OptionsAttr) -> proc_macro2::TokenStream
 {
     let engine_crate_path = find_engine_crate_path().unwrap();
 
-    let input_ident = &input.ident;
-
-    let mut generics = input.generics;
-
-    for type_param in generics.type_params_mut() {
-        type_param
-            .bounds
-            .push(syn::TypeParamBound::Trait(syn::TraitBound {
-                paren_token: None,
-                modifier: syn::TraitBoundModifier::None,
-                lifetimes: None,
-                path: engine_crate_path.join(syn_path!(reflection::Reflection)),
-            }));
-    }
-
-    let (impl_generics, type_generics, where_clause) = generics.split_for_impl();
-
-    let variant_lookup_match_arms =
-        input.variants.iter().enumerate().map(|(index, variant)| {
+    let variant_lookup_match_arms = input
+        .variants
+        .iter()
+        .enumerate()
+        .map(|(index, variant)| {
             let variant_ident = &variant.ident;
 
             let pattern = match variant.fields {
@@ -43,178 +26,74 @@ pub fn generate(input: syn::ItemEnum) -> proc_macro2::TokenStream
             quote! {
                 #pattern => &enum_reflection.variants[#index]
             }
-        });
+        })
+        .collect::<Vec<_>>();
 
     let is_unit_only = input
         .variants
         .iter()
         .all(|variant| matches!(variant.fields, syn::Fields::Unit));
 
-    let mod_name = format_ident!("__engine_private_{input_ident}");
-
-    let mut has_c_repr = false;
-
-    let mut primitive_repr: Option<proc_macro2::Ident> = None;
-
-    for attr in &input.attrs {
-        let syn::Meta::List(attr_meta) = &attr.meta else {
-            continue;
-        };
-
-        if !attr_meta.path.is_ident("repr") {
-            continue;
-        }
-
-        attr_meta
-            .parse_nested_meta(|nested_meta| {
-                let Some(meta_ident) = nested_meta.path.get_ident() else {
-                    return Ok(());
-                };
-
-                if meta_ident == "C" {
-                    has_c_repr = true;
-                } else if PRIMITIVE_REPRS.contains(&meta_ident.to_string().as_str()) {
-                    primitive_repr = Some(meta_ident.clone());
-                } else {
-                    return Err(nested_meta.error(concat!(
-                        "Unsupported representation. ",
-                        "Only C and primitive representations are allowed"
-                    )));
-                }
+    let mod_name = format_ident!("__engine_private_{}", input.ident);
 
-                Ok(())
-            })
-            .unwrap();
-    }
+    let reprs = get_reprs(&input);
 
-    if !has_c_repr && primitive_repr.is_none() {
+    if !reprs.has_c_repr && reprs.primitive_repr.is_none() {
         panic!("Enums must have a C or primitive representation to derive Reflection");
     }
 
-    let mod_content = if is_unit_only {
-        quote! {}
-    } else {
-        let variant_fields_structs =
-            input.variants.iter().map(|variant| {
-                let fields = variant.fields.iter().map(|field| {
-                    let field_type = &field.ty;
-
-                    if let Some(field_ident) = &field.ident {
-                        quote! { pub #field_ident: #field_type }
-                    } else {
-                        quote! { pub #field_type }
-                    }
-                });
-
-                let ident = format_ident!("VariantFields{}", variant.ident);
-
-                let generics_phantom_data_elems = generics.params.iter().filter_map(
-                    |generic_param| match generic_param {
-                        syn::GenericParam::Type(type_param) => {
-                            Some(type_param.ident.clone().into_token_stream())
-                        }
-                        syn::GenericParam::Lifetime(_) => {
-                            unimplemented!();
-                        }
-                        syn::GenericParam::Const(_const_param) => None,
-                    },
-                );
-
-                let generics_phantom_data = quote! {
-                    std::marker::PhantomData<(#(#generics_phantom_data_elems),*)>
-                };
-
-                match variant.fields {
-                    syn::Fields::Named(_) => quote! {
-                        #[repr(C)]
-                        pub struct #ident #impl_generics
-                            #where_clause
-                        {
-                            #(#fields,)*
-                            _pd: #generics_phantom_data
-                        }
-                    },
-                    syn::Fields::Unnamed(_) => quote! {
-                        #[repr(C)]
-                        pub struct #ident #impl_generics (
-                            #(#fields,)*
-                            #generics_phantom_data
-                        )
-                            #where_clause;
-                    },
-                    syn::Fields::Unit => quote! {
-                        #[repr(C)]
-                        pub struct #ident #impl_generics (#generics_phantom_data)
-                            #where_clause;
-                    },
-                }
-            });
-
-        let fields_union_fields = input.variants.iter().map(|variant| {
-            let variant_ident = &variant.ident;
-
-            let variant_fields_struct_ident =
-                format_ident!("VariantFields{}", variant.ident);
-
-            quote! {
-                pub #variant_ident:
-                    std::mem::ManuallyDrop<#variant_fields_struct_ident #type_generics>
-            }
-        });
-
-        let discriminant_enum_variants = input.variants.iter().map(|variant| {
-            let variant_ident = &variant.ident;
-
-            if let Some((_, discriminant)) = &variant.discriminant {
-                quote! { #variant_ident = #discriminant }
-            } else {
-                quote! { #variant_ident }
-            }
-        });
-
-        // If the enum has both a C & primitive repr, the primitive repr must be used
-        let discriminant_enum_repr = if let Some(primitive_repr) = primitive_repr {
-            quote! { #primitive_repr }
-        } else if has_c_repr {
-            quote! { C }
+    let mod_content = generate_mod_content(&input, is_unit_only, &reprs);
+
+    let impls: &mut dyn Iterator<Item = proc_macro2::TokenStream> =
+        if input.generics.params.is_empty() {
+            &mut [generate_impls(
+                &input,
+                None,
+                is_unit_only,
+                &variant_lookup_match_arms,
+                &mod_name,
+                &engine_crate_path,
+            )]
+            .into_iter()
         } else {
-            unreachable!();
+            &mut options.impl_with_generics.iter().map(|generic_args| {
+                generate_impls(
+                    &input,
+                    Some(generic_args),
+                    is_unit_only,
+                    &variant_lookup_match_arms,
+                    &mod_name,
+                    &engine_crate_path,
+                )
+            })
         };
 
-        quote! {
-            #![allow(non_snake_case, dead_code)]
-
-            use super::*;
-
-            #[repr(C)]
-            pub struct Equivalent #impl_generics
-                #where_clause
-            {
-                pub tag: Discriminant,
-                pub payload: Fields #type_generics
-            }
-
-            #[repr(#discriminant_enum_repr)]
-            pub enum Discriminant
-            {
-                #(#discriminant_enum_variants),*
-            }
+    quote! {
+        #[doc(hidden)]
+        #[allow(non_snake_case)]
+        mod #mod_name {
+            #mod_content
+        }
 
-            #[repr(C)]
-            pub union Fields #impl_generics
-                #where_clause
-            {
-                #(#fields_union_fields),*
-            }
+        #(#impls)*
+    }
+}
 
-            #(#variant_fields_structs)*
-        }
-    };
+fn generate_impls(
+    input: &syn::ItemEnum,
+    generic_args: Option<&syn::AngleBracketedGenericArguments>,
+    is_unit_only: bool,
+    variant_lookup_match_arms: &[proc_macro2::TokenStream],
+    mod_name: &proc_macro2::Ident,
+    engine_crate_path: &syn::Path,
+) -> proc_macro2::TokenStream
+{
+    let get_default_value_fn = gen_get_default_value_fn(&input.ident, generic_args);
 
     let variants = generate_variants(
-        input_ident,
+        &input.ident,
         &input.variants,
-        &generics,
+        generic_args,
         is_unit_only,
         &engine_crate_path,
     );
@@ -228,28 +107,46 @@ pub fn generate(input: syn::ItemEnum) -> proc_macro2::TokenStream
             Some(#engine_crate_path::reflection::EnumTaggedUnion {
                 discriminant_layout: std::alloc::Layout::new::<#mod_name::Discriminant>(),
                 discriminant_byte_offset:
-                    std::mem::offset_of!(#mod_name::Equivalent #type_generics, tag),
+                    std::mem::offset_of!(#mod_name::Equivalent #generic_args, tag),
                 fields_layout:
-                    std::alloc::Layout::new::<#mod_name::Fields #type_generics>(),
+                    std::alloc::Layout::new::<#mod_name::Fields #generic_args>(),
                 fields_byte_offset:
-                    std::mem::offset_of!(#mod_name::Equivalent #type_generics, payload),
+                    std::mem::offset_of!(#mod_name::Equivalent #generic_args, payload),
             })
         }
     };
 
-    let get_default_value_fn = gen_get_default_value_fn(&input_ident, &type_generics);
+    let generics_type_aliases = input
+        .generics
+        .type_params()
+        .zip(
+            generic_args
+                .iter()
+                .map(|generic_args| &generic_args.args)
+                .flatten()
+                .flat_map(|generic_arg| match generic_arg {
+                    syn::GenericArgument::Type(ty) => Some(ty),
+                    _ => None,
+                }),
+        )
+        .map(|(type_param, generic_arg_type)| {
+            let type_param_ident = &type_param.ident;
 
-    quote! {
-        #[doc(hidden)]
-        mod #mod_name {
-            #mod_content
-        }
+            quote! {
+                type #type_param_ident = #generic_arg_type;
+            }
+        });
+
+    let input_ident = &input.ident;
 
-        unsafe impl #impl_generics #engine_crate_path::reflection::Reflection for
-            #input_ident #type_generics #where_clause
+    quote! {
+        unsafe impl #engine_crate_path::reflection::Reflection for
+            #input_ident #generic_args
         {
             const TYPE_REFLECTION: &#engine_crate_path::reflection::Type =
                 &const {
+                    #(#generics_type_aliases)*
+
                     #engine_crate_path::reflection::Type::Enum(
                         #engine_crate_path::reflection::Enum {
                             variants: &[#(#variants),*],
@@ -263,8 +160,8 @@ pub fn generate(input: syn::ItemEnum) -> proc_macro2::TokenStream
                 };
         }
 
-        unsafe impl #impl_generics #engine_crate_path::reflection::EnumReflectionExt for
-            #input_ident #type_generics #where_clause
+        unsafe impl #engine_crate_path::reflection::EnumReflectionExt for
+            #input_ident #generic_args
         {
             fn get_variant_reflection(&self)
                 -> &'static #engine_crate_path::reflection::EnumVariant
@@ -286,15 +183,13 @@ pub fn generate(input: syn::ItemEnum) -> proc_macro2::TokenStream
 fn generate_variants<'a>(
     input_ident: &proc_macro2::Ident,
     input_variants: &'a syn::punctuated::Punctuated<syn::Variant, syn::Token![,]>,
-    input_generics: &'a syn::Generics,
+    generic_args: Option<&'a syn::AngleBracketedGenericArguments>,
     is_unit_only: bool,
     engine_crate_path: &'a syn::Path,
 ) -> impl Iterator<Item = proc_macro2::TokenStream> + use<'a>
 {
     let mod_name = format_ident!("__engine_private_{input_ident}");
 
-    let (_, type_generics, _) = input_generics.split_for_impl();
-
     input_variants.iter().map(move |variant| {
         let variant_name =
             syn::LitStr::new(&variant.ident.to_string(), variant.ident.span());
@@ -321,12 +216,12 @@ fn generate_variants<'a>(
                                     if let Some(field_ident) = &field.ident {
                                         quote! {
                                             std::mem::offset_of!(
-                                                #mod_name::Equivalent #type_generics,
+                                                #mod_name::Equivalent #generic_args,
                                                 payload.#variant_ident
                                             ) +
                                             std::mem::offset_of!(
                                                 #mod_name::#variant_fields_struct_ident
-                                                    #type_generics,
+                                                    #generic_args,
                                                 #field_ident
                                             )
                                         }
@@ -334,9 +229,6 @@ fn generate_variants<'a>(
                                         unreachable!();
                                     }
                                 },
-                                type_reflection_optional: input_generics
-                                    .params
-                                    .is_empty(),
                             },
                         )
                     },
@@ -371,20 +263,17 @@ fn generate_variants<'a>(
 
                                         quote! {
                                             std::mem::offset_of!(
-                                                #mod_name::Equivalent #type_generics,
+                                                #mod_name::Equivalent #generic_args,
                                                 payload.#variant_ident
                                             ) +
                                             std::mem::offset_of!(
                                                 #mod_name::#variant_fields_struct_ident
-                                                    #type_generics,
+                                                    #generic_args,
                                                 #field_index
                                             )
                                         }
                                     }
                                 },
-                                type_reflection_optional: input_generics
-                                    .params
-                                    .is_empty(),
                             },
                         )
                     },
@@ -449,3 +338,181 @@ fn generate_variants<'a>(
         }
     })
 }
+
+fn generate_mod_content(
+    input: &syn::ItemEnum,
+    is_unit_only: bool,
+    reprs: &Reprs,
+) -> proc_macro2::TokenStream
+{
+    if is_unit_only {
+        return quote! {};
+    }
+
+    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
+
+    let variant_fields_structs = input.variants.iter().map(|variant| {
+        let fields = variant.fields.iter().map(|field| {
+            let field_type = &field.ty;
+
+            if let Some(field_ident) = &field.ident {
+                quote! { pub #field_ident: #field_type }
+            } else {
+                quote! { pub #field_type }
+            }
+        });
+
+        let ident = format_ident!("VariantFields{}", variant.ident);
+
+        let generics_phantom_data_elems = input.generics.params.iter().filter_map(
+            |generic_param| match generic_param {
+                syn::GenericParam::Type(type_param) => {
+                    Some(type_param.ident.clone().into_token_stream())
+                }
+                syn::GenericParam::Lifetime(_) => {
+                    unimplemented!();
+                }
+                syn::GenericParam::Const(_const_param) => None,
+            },
+        );
+
+        let generics_phantom_data = quote! {
+            std::marker::PhantomData<(#(#generics_phantom_data_elems),*)>
+        };
+
+        match variant.fields {
+            syn::Fields::Named(_) => quote! {
+                #[repr(C)]
+                pub struct #ident #impl_generics
+                    #where_clause
+                {
+                    #(#fields,)*
+                    _pd: #generics_phantom_data
+                }
+            },
+            syn::Fields::Unnamed(_) => quote! {
+                #[repr(C)]
+                pub struct #ident #impl_generics (
+                    #(#fields,)*
+                    #generics_phantom_data
+                )
+                    #where_clause;
+            },
+            syn::Fields::Unit => quote! {
+                #[repr(C)]
+                pub struct #ident #impl_generics (#generics_phantom_data)
+                    #where_clause;
+            },
+        }
+    });
+
+    let fields_union_fields = input.variants.iter().map(|variant| {
+        let variant_ident = &variant.ident;
+
+        let variant_fields_struct_ident = format_ident!("VariantFields{}", variant.ident);
+
+        quote! {
+            pub #variant_ident:
+                std::mem::ManuallyDrop<#variant_fields_struct_ident #type_generics>
+        }
+    });
+
+    let discriminant_enum_variants = input.variants.iter().map(|variant| {
+        let variant_ident = &variant.ident;
+
+        if let Some((_, discriminant)) = &variant.discriminant {
+            quote! { #variant_ident = #discriminant }
+        } else {
+            quote! { #variant_ident }
+        }
+    });
+
+    // If the enum has both a C & primitive repr, the primitive repr must be used
+    let discriminant_enum_repr = if let Some(primitive_repr) = &reprs.primitive_repr {
+        quote! { #primitive_repr }
+    } else if reprs.has_c_repr {
+        quote! { C }
+    } else {
+        unreachable!();
+    };
+
+    quote! {
+        #![allow(non_snake_case, dead_code)]
+
+        use super::*;
+
+        #[repr(C)]
+        pub struct Equivalent #impl_generics
+            #where_clause
+        {
+            pub tag: Discriminant,
+            pub payload: Fields #type_generics
+        }
+
+        #[repr(#discriminant_enum_repr)]
+        pub enum Discriminant
+        {
+            #(#discriminant_enum_variants),*
+        }
+
+        #[repr(C)]
+        pub union Fields #impl_generics
+            #where_clause
+        {
+            #(#fields_union_fields),*
+        }
+
+        #(#variant_fields_structs)*
+    }
+}
+
+fn get_reprs(input: &ItemEnum) -> Reprs
+{
+    let mut has_c_repr = false;
+    let mut primitive_repr: Option<proc_macro2::Ident> = None;
+
+    for attr in &input.attrs {
+        let syn::Meta::List(attr_meta) = &attr.meta else {
+            continue;
+        };
+
+        if !attr_meta.path.is_ident("repr") {
+            continue;
+        }
+
+        attr_meta
+            .parse_nested_meta(|nested_meta| {
+                let Some(meta_ident) = nested_meta.path.get_ident() else {
+                    return Ok(());
+                };
+
+                if meta_ident == "C" {
+                    has_c_repr = true;
+                } else if PRIMITIVE_REPRS.contains(&meta_ident.to_string().as_str()) {
+                    primitive_repr = Some(meta_ident.clone());
+                } else {
+                    return Err(nested_meta.error(concat!(
+                        "Unsupported representation. ",
+                        "Only C and primitive representations are allowed"
+                    )));
+                }
+
+                Ok(())
+            })
+            .unwrap();
+    }
+
+    Reprs { has_c_repr, primitive_repr }
+}
+
+#[derive(Debug)]
+struct Reprs
+{
+    has_c_repr: bool,
+    primitive_repr: Option<proc_macro2::Ident>,
+}
+
+const PRIMITIVE_REPRS: &[&'static str] = &[
+    "u8", "u16", "u32", "u64", "u128", "usize", "i8", "i16", "i32", "i64", "i128",
+    "isize",
+];