feat(ecs-macros): add support for setting component ref types

1 files changed, 65 insertions, 2 deletions

diff --git a/ecs-macros/src/lib.rs b/ecs-macros/src/lib.rs
index 8bfce38..8bef0b6 100644
--- a/ecs-macros/src/lib.rs
+++ b/ecs-macros/src/lib.rs
@@ -14,6 +14,7 @@ use syn::{
     ItemStruct,
     ItemUnion,
     Path,
+    Type,
 };
 use toml::value::{Table as TomlTable, Value as TomlValue};
 
@@ -46,6 +47,10 @@ pub fn component_derive(input: TokenStream) -> TokenStream
 {
     let item: TypeItem = parse::<Item>(input).unwrap().try_into().unwrap();
 
+    let ComponentAttribute { ref_type, ref_mut_type } = item
+        .attribute::<ComponentAttribute>("component")
+        .unwrap_or_default();
+
     let item_ident = item.ident();
 
     let (impl_generics, type_generics, where_clause) = item.generics().split_for_impl();
@@ -120,8 +125,8 @@ pub fn component_derive(input: TokenStream) -> TokenStream
             {
                 type Component = Self;
 
-                type RefMut<'component> = ComponentRefMut<'component, Self>;
-                type Ref<'component> = ComponentRef<'component, Self>;
+                type RefMut<'component> = #ref_mut_type;
+                type Ref<'component> = #ref_type;
 
                 fn id() -> Uid
                 {
@@ -369,3 +374,61 @@ fn find_engine_ecs_crate_path() -> Option<Path>
         None
     })
 }
+
+#[derive(Debug)]
+struct ComponentAttribute
+{
+    ref_type: proc_macro2::TokenStream,
+    ref_mut_type: proc_macro2::TokenStream,
+}
+
+impl FromAttribute for ComponentAttribute
+{
+    fn from_attribute(attribute: &Attribute) -> Result<Self, syn::Error>
+    {
+        let mut ref_type: Option<Type> = None;
+        let mut ref_mut_type: Option<Type> = None;
+
+        attribute.parse_nested_meta(|meta| {
+            let Some(flag) = meta.path.get_ident() else {
+                return Err(meta.error("Not a single identifier"));
+            };
+
+            if flag == "ref_type" {
+                let value = meta.value()?;
+
+                ref_type = Some(value.parse::<Type>()?);
+
+                return Ok(());
+            } else if flag == "ref_mut_type" {
+                let value = meta.value()?;
+
+                ref_mut_type = Some(value.parse::<Type>()?);
+
+                return Ok(());
+            }
+
+            Err(meta.error("Unrecognized token"))
+        })?;
+
+        Ok(Self {
+            ref_type: ref_type
+                .map(|ref_type| ref_type.into_token_stream())
+                .unwrap_or_else(|| Self::default().ref_type),
+            ref_mut_type: ref_mut_type
+                .map(|ref_mut_type| ref_mut_type.into_token_stream())
+                .unwrap_or_else(|| Self::default().ref_mut_type),
+        })
+    }
+}
+
+impl Default for ComponentAttribute
+{
+    fn default() -> Self
+    {
+        Self {
+            ref_type: quote! { ComponentRef<'component, Self> },
+            ref_mut_type: quote! { ComponentRefMut<'component, Self> },
+        }
+    }
+}