feat(engine): add enum reflection

1 files changed, 78 insertions, 5 deletions

diff --git a/engine-macros/src/lib.rs b/engine-macros/src/lib.rs
index 75ff209..a508d8a 100644
--- a/engine-macros/src/lib.rs
+++ b/engine-macros/src/lib.rs
@@ -2,7 +2,7 @@
 
 use proc_macro::TokenStream;
 use quote::{ToTokens, format_ident, quote};
-use syn::{Item, LitStr, Path as SynPath, parse};
+use syn::{Fields, Item, ItemEnum, ItemStruct, LitStr, Path as SynPath, parse};
 
 macro_rules! syn_path {
     ($first_segment: ident $(::$segment: ident)*) => {
@@ -33,12 +33,15 @@ pub fn reflection_derive(input: TokenStream) -> TokenStream
 {
     let input = parse::<Item>(input).unwrap();
 
-    let input = match input {
-        Item::Struct(input) => input,
-        Item::Enum(_) => unimplemented!(),
+    match input {
+        Item::Struct(input) => generate_struct_reflection_impl(input),
+        Item::Enum(input) => generate_enum_reflection_impl(input),
         _ => panic!("Invalid input"),
-    };
+    }
+}
 
+fn generate_struct_reflection_impl(input: ItemStruct) -> TokenStream
+{
     let engine_crate_path = find_engine_crate_path().unwrap();
 
     let input_ident = input.ident;
@@ -124,6 +127,76 @@ pub fn reflection_derive(input: TokenStream) -> TokenStream
     .into()
 }
 
+fn generate_enum_reflection_impl(input: ItemEnum) -> TokenStream
+{
+    let engine_crate_path = find_engine_crate_path().unwrap();
+
+    let input_ident = input.ident;
+
+    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
+
+    let variants = input.variants.iter().map(|variant| {
+        let variant_name = LitStr::new(&variant.ident.to_string(), variant.ident.span());
+
+        quote! {
+            #engine_crate_path::reflection::EnumVariant {
+                name: #variant_name,
+            }
+        }
+    });
+
+    let variant_lookup_match_arms =
+        input.variants.iter().enumerate().map(|(index, variant)| {
+            let variant_ident = &variant.ident;
+
+            let pattern = match variant.fields {
+                Fields::Unit => quote! { Self::#variant_ident },
+                Fields::Named(_) => quote! { Self::#variant_ident { .. } },
+                Fields::Unnamed(_) => quote! { Self::#variant_ident(..) },
+            };
+
+            quote! {
+                #pattern => &enum_reflection.variants[#index]
+            }
+        });
+
+    quote! {
+        unsafe impl #impl_generics #engine_crate_path::reflection::Reflection for
+            #input_ident #type_generics #where_clause
+        {
+            const TYPE_REFLECTION: &#engine_crate_path::reflection::Type =
+                &const {
+                    #engine_crate_path::reflection::Type::Enum(
+                        #engine_crate_path::reflection::Enum {
+                            variants: &[
+                                #(#variants),*
+                            ]
+                        }
+                    )
+                };
+        }
+
+        unsafe impl #impl_generics #engine_crate_path::reflection::EnumReflectionExt for
+            #input_ident #type_generics #where_clause
+        {
+            fn get_variant_reflection(&self)
+                -> &'static #engine_crate_path::reflection::EnumVariant
+            {
+                let enum_reflection = unsafe {
+                    <Self as #engine_crate_path::reflection::Reflection>::TYPE_REFLECTION
+                        .as_enum()
+                        .unwrap_unchecked()
+                };
+
+                match self {
+                    #(#variant_lookup_match_arms),*
+                }
+            }
+        }
+    }
+    .into()
+}
+
 fn find_engine_crate_path() -> Option<SynPath>
 {
     let cargo_crate_name = std::env::var("CARGO_CRATE_NAME").ok()?;