feat(engine): add try_get_field_mut fn to enum variant reflection

2 files changed, 116 insertions, 2 deletions

diff --git a/engine-macros/src/reflection/enum_impl.rs b/engine-macros/src/reflection/enum_impl.rs
index 4e27c2b..b637d9f 100644
--- a/engine-macros/src/reflection/enum_impl.rs
+++ b/engine-macros/src/reflection/enum_impl.rs
@@ -273,6 +273,61 @@ fn generate_variants<'a>(
                 }
             };
 
+        let variant_ident = &variant.ident;
+
+        let variant_pattern = match &variant.fields {
+            syn::Fields::Unit => quote! { Self::#variant_ident },
+            syn::Fields::Named(syn::FieldsNamed { named: named_fields, .. }) => {
+                let field_idents = named_fields.iter().map(|field| {
+                    let Some(field_ident) = &field.ident else {
+                        unreachable!();
+                    };
+
+                    field_ident
+                });
+
+                quote! { Self::#variant_ident { #(#field_idents),* } }
+            }
+            syn::Fields::Unnamed(syn::FieldsUnnamed {
+                unnamed: unnamed_fields, ..
+            }) => {
+                let field_idents = unnamed_fields
+                    .iter()
+                    .enumerate()
+                    .map(|(field_index, _)| format_ident!("field_{field_index}"));
+
+                quote! { Self::#variant_ident(#(#field_idents),*) }
+            }
+        };
+
+        let field_index_match_arms = match &variant.fields {
+            syn::Fields::Unit => quote! {},
+            syn::Fields::Named(syn::FieldsNamed { named: named_fields, .. }) => {
+                let match_arms =
+                    named_fields.iter().enumerate().map(|(field_index, field)| {
+                        let Some(field_ident) = &field.ident else {
+                            unreachable!();
+                        };
+
+                        quote! { #field_index => { #field_ident } }
+                    });
+
+                quote! { #(#match_arms)* }
+            }
+            syn::Fields::Unnamed(syn::FieldsUnnamed {
+                unnamed: unnamed_fields, ..
+            }) => {
+                let match_arms =
+                    unnamed_fields.iter().enumerate().map(|(field_index, _)| {
+                        let field_ident = format_ident!("field_{field_index}");
+
+                        quote! { #field_index => { #field_ident } }
+                    });
+
+                quote! { #(#match_arms)* }
+            }
+        };
+
         quote! {
             #engine_crate_path::reflection::EnumVariant {
                 name: #variant_name,
@@ -296,7 +351,30 @@ fn generate_variants<'a>(
                     *dst = created;
 
                     Ok(())
-                }
+                },
+                try_get_field_mut: |target, field_index| { {
+                    #![allow(unreachable_code)]
+
+                    use std::any::Any;
+                    use #engine_crate_path::reflection::EnumVariantGetFieldError;
+
+                    let target = target
+                        .downcast_mut::<Self>()
+                        .ok_or(EnumVariantGetFieldError::WrongTargetType)?;
+
+                    let #variant_pattern = target else {
+                        return Err(EnumVariantGetFieldError::WrongTargetVariant);
+                    };
+
+                    let field: &mut dyn Any = match field_index {
+                        #field_index_match_arms
+                        _ => {
+                            return Err(EnumVariantGetFieldError::FieldIndexOutOfBounds);
+                        }
+                    };
+
+                    Ok(field)
+                } },
             }
         }
     })
diff --git a/engine-reflection/src/lib.rs b/engine-reflection/src/lib.rs
index a31fef0..79a7417 100644
--- a/engine-reflection/src/lib.rs
+++ b/engine-reflection/src/lib.rs
@@ -168,12 +168,17 @@ pub struct EnumVariant
         dst: &mut dyn Any,
         fields: &mut dyn Iterator<Item = Box<dyn Any>>,
     ) -> Result<(), EnumVariantWriteNewToError>,
+
+    pub try_get_field_mut: fn(
+        target: &mut dyn Any,
+        field_index: usize,
+    ) -> Result<&mut dyn Any, EnumVariantGetFieldError>,
 }
 
 #[derive(Debug, thiserror::Error)]
 pub enum EnumVariantWriteNewToError
 {
-    #[error("Destination must be the same type as the variant's enum")]
+    #[error("Destination must have the same type as the variant's enum")]
     WrongDstType,
 
     #[error("Too few fields are specified")]
@@ -186,6 +191,19 @@ pub enum EnumVariantWriteNewToError
     WrongFieldType,
 }
 
+#[derive(Debug, thiserror::Error)]
+pub enum EnumVariantGetFieldError
+{
+    #[error("Target is not the same type as the variant's enum")]
+    WrongTargetType,
+
+    #[error("Target is not the correct variant")]
+    WrongTargetVariant,
+
+    #[error("Field index is out of bounds")]
+    FieldIndexOutOfBounds,
+}
+
 #[derive(Debug, Clone)]
 pub enum EnumVariantFields
 {
@@ -468,6 +486,21 @@ unsafe impl<T: Reflection> Reflection for Option<T>
 
                     Ok(())
                 },
+                try_get_field_mut: |target, field_index| {
+                    if field_index != 0 {
+                        return Err(EnumVariantGetFieldError::FieldIndexOutOfBounds);
+                    }
+
+                    let target = target
+                        .downcast_mut::<Self>()
+                        .ok_or(EnumVariantGetFieldError::WrongTargetType)?;
+
+                    let Some(value_field) = target else {
+                        return Err(EnumVariantGetFieldError::WrongTargetVariant);
+                    };
+
+                    Ok(value_field)
+                },
             },
             EnumVariant {
                 name: "None",
@@ -485,6 +518,9 @@ unsafe impl<T: Reflection> Reflection for Option<T>
 
                     Ok(())
                 },
+                try_get_field_mut: |_, _| {
+                    Err(EnumVariantGetFieldError::FieldIndexOutOfBounds)
+                },
             },
         ],
         is_unit_only: false,