lyra-engine/lyra-reflect/lyra-reflect-derive/src/struct_derive.rs

use quote::{quote, ToTokens};
use syn::{DeriveInput, parse_quote, DataStruct};

use crate::{add_trait_bounds, FieldAttributes};

#[derive(Debug, PartialEq, Eq)]
enum StructType {
    Unit,
    Named,
    Tuple,
}

impl StructType {
    pub fn new(data: &DataStruct) -> Self {
        if let Some(first) = data.fields.iter().next() {
            if first.ident.is_some() {
                Self::Named
            } else {
                Self::Tuple
            }
        } else {
            Self::Unit
        }
    }
}

/// Generates code that matches a string with a struct's field name, and returns an Option that
/// contains a borrow (mutable borrow if `is_mut` is true) to the matching struct field.
/// 
/// Example:
/// ```compile_fail
/// // when `is_mut` = false
/// match name {
///     "x" => Some(&self.x),
///     "y" => Some(&self.y),
///     _ => None,
/// }
/// 
/// // when `is_mut` = true
/// match name {
///     "x" => Some(&mut self.x),
///     "y" => Some(&mut self.y),
///     _ => None,
/// }
/// ```
/// 
/// If the struct is a unit or tuple struct, None will always be returned
fn gen_struct_field_match(data: &DataStruct, is_mut: bool) -> proc_macro2::TokenStream {
    let ty = StructType::new(data);

    if ty == StructType::Named {
        let mut_tkn = if is_mut {
            quote! {
                mut
            }
        } else { quote!{} };

        let field_arms = data.fields.iter().map(|field| {
            let field_ident = field.ident.as_ref().expect("Struct is missing field ident!");
            let field_name_str = field_ident.to_string();
            
            let attrs = FieldAttributes::from_vec(&field.attrs)
                .expect("Failure to parse reflect attributes");
            if attrs.has_skip() {
                quote! {
                    #field_name_str => None
                }
            } else {
                quote! {
                    #field_name_str => Some(&#mut_tkn self.#field_ident)
                }
            }
        });

        quote! {
            match name {
                #(#field_arms,)*
                _ => None,
            }
        }
    } else { quote!(None) }
}

/// Generates code that matches a string with a struct's field name, and sets that field value
/// with the provided `val`.
/// 
/// Example:
/// ```compile_fail
/// match name {
///     "x" => self.x = any_val.downcast_ref::<f32>()
///                .expect(&format!("Cannot set struct's field of {} type to the provided type of {}", "f32", val.name()))
///                .clone(),
///     "y" => self.y = any_val.downcast_ref::<f32>()
///                .expect(&format!("Cannot set struct's field of {} type to the provided type of {}", "f32", val.name()))
///                .clone()
///     _ => {
///         return false;
///     },
/// }
/// ```
fn gen_struct_set_field_match(data: &DataStruct) -> proc_macro2::TokenStream {
    let ty = StructType::new(data);

    if ty == StructType::Named {
        let field_arms = data.fields.iter().map(|field| {
            let field_ident = field.ident.as_ref().expect("Struct is missing field ident!");
            let field_name_str = field_ident.to_string();
            let field_ty = &field.ty;
            
            let attrs = FieldAttributes::from_vec(&field.attrs)
                .expect("Failure to parse reflect attributes");
            if attrs.has_skip() {
                quote! {
                    #field_name_str => {}
                }
            } else {
                quote! {
                    #field_name_str => self.#field_ident = any_val.downcast_ref::<#field_ty>()
                        .expect(&format!("Cannot set struct's field of {} type to the provided type of {}", #field_name_str, val.name()))
                        .clone() //Some(&#mut_tkn self.#field_ident)
                }
            }
        });

        quote! {
            let any_val = val.as_any();
            match name {
                #(#field_arms,)*
                _ => {
                    return false;
                },
            }

            true
        }
    } else { quote!(return false;) }
}

/// Generates code that matches a string with a struct's field name, and returns a string that is
/// the type of the field.
/// 
/// Example:
/// ```compile_fail
/// match name {
///     "x" => Some("f32"),
///     "y" => Some("f32"),
///     _ => None,
/// }
/// ```
fn gen_struct_field_name_match(data: &DataStruct) -> proc_macro2::TokenStream {
    let ty = StructType::new(data);

    if ty == StructType::Named {
        let field_arms = data.fields.iter().map(|field| {
            let field_ident = field.ident.as_ref().expect("Struct is missing field ident!");
            let field_name_str = field_ident.to_string();

            let mut field_ty_stream = proc_macro2::TokenStream::new();
            field.ty.to_tokens(&mut field_ty_stream);
            let s = field_ty_stream.to_string();
            
            quote! {
                #field_name_str => Some(#s)
            }
        });

        quote! {
            match name {
                #(#field_arms,)*
                _ => None,
            }
        }
    } else { quote!(None) }
}

/// Generates code that matches a string with a struct's field name, and sets that field value
/// with the provided `val`.
/// 
/// Example:
/// ```compile_fail
/// match name {
///     0 => self.x = any_val.downcast_ref::<f32>()
///                .expect(&format!("Cannot set struct's field of {} type to the provided type of {}", "f32", val.name()))
///                .clone(),
///     1 => self.y = any_val.downcast_ref::<f32>()
///                .expect(&format!("Cannot set struct's field of {} type to the provided type of {}", "f32", val.name()))
///                .clone()
///     _ => {
///         return false;
///     },
/// }
/// ```
fn gen_struct_set_field_match_idx(data: &DataStruct) -> proc_macro2::TokenStream {
    let ty = StructType::new(data);

    if ty == StructType::Unit {
        return quote!( return false; );
    }

    let field_arms = data.fields.iter().enumerate().map(|(idx, field)| {
        let field_ty = &field.ty;

        let attrs = FieldAttributes::from_vec(&field.attrs)
            .expect("Failure to parse reflect attributes");
        if attrs.has_skip() {
            return quote! {
                #idx => {}
            };
        }

        if let Some(field_ident) = &field.ident {
            let field_name_str = field_ident.to_string();

            quote! {
                #idx => self.#field_ident = any_val.downcast_ref::<#field_ty>()
                    .expect(&format!("Cannot set struct's field of {} to the provided type of {}", #field_name_str, val.name()))
                    .clone()
            }
        } else {
            let sidx = syn::Index::from(idx);
            quote! {
                #idx => self.#sidx = any_val.downcast_ref::<#field_ty>()
                    .expect(&format!("Cannot set struct's field at {} to the provided type of {}", #idx, val.name()))
                    .clone()
            }
        }
    });

    quote! {
        let any_val = val.as_any();
        match idx {
            #(#field_arms,)*
            _ => {
                return false;
            },
        }

        true
    }
}

/// Generates code that matches an index with the correct field in the struct, and returns the
/// type of the field.
/// 
/// Example:
/// ```compile_fail
/// match name {
///     0 => Some("f32"),
///     1 => Some("f32"),
///     _ => None,
/// }
/// ```
fn gen_struct_field_name_match_idx(data: &DataStruct) -> proc_macro2::TokenStream {
    let field_arms = data.fields.iter().enumerate().map(|(idx, field)| {
        let mut field_ty_stream = proc_macro2::TokenStream::new();
        field.ty.to_tokens(&mut field_ty_stream);
        let field_ty_str = field_ty_stream.to_string();
        
        quote! {
            #idx => Some(#field_ty_str)
        }
    });

    quote! {
        match idx {
            #(#field_arms,)*
            _ => None,
        }
    }
}

/// Generates code that matches an index with the position of the field in a struct,
/// and returns an Option that contains a borrow (mutable borrow if `is_mut` is true)
/// to the matching struct field.
/// 
/// Example:
/// ```compile_fail
/// // when `is_mut` = false
/// match idx {
///     0 => Some(&self.x),
///     1 => Some(&self.y),
///     _ => None,
/// }
/// 
/// // when `is_mut` = true
/// match idx {
///     0 => Some(&mut self.x),
///     1 => Some(&mut self.y),
///     _ => None,
/// }
/// ```
fn gen_struct_field_match_idx(data: &DataStruct, is_mut: bool) -> proc_macro2::TokenStream {
    let ty = StructType::new(data);

    if ty == StructType::Unit {
        return quote!(None);
    }

    let mut_tkn = if is_mut {
        quote! {
            mut
        }
    } else { quote!{} };
    
    let field_arms = data.fields.iter().enumerate().map(|(idx, field)| {
        let attrs = FieldAttributes::from_vec(&field.attrs)
            .expect("Failure to parse reflect attributes");
        if attrs.has_skip() {
            return quote! {
                #idx => None
            };
        }
        
        if let Some(field_ident) = &field.ident {
            quote! {
                #idx => Some(&#mut_tkn self.#field_ident)
            }
        } else {
            let sidx = syn::Index::from(idx);
            quote! {
                #idx => Some(&#mut_tkn self.#sidx)
            }
        }
    });

    quote! {
        match idx {
            #(#field_arms,)*
            _ => None,
        }
    }
}

/// Generates code that matches an index with the position of the field in a struct,
/// and returns an Option that contains the name of the field.
/// 
/// Example:
/// ```compile_fail
/// match idx {
///     0 => Some("x"),
///     1 => Some("y"),
///     _ => None,
/// }
/// ```
fn gen_struct_field_name_idx(data: &DataStruct) -> proc_macro2::TokenStream {
    let ty = StructType::new(data);

    if ty == StructType::Named {
        let field_arms = data.fields.iter().enumerate().map(|(idx, field)| {
            let field_ident = field.ident.as_ref().expect("Struct is missing field ident!");
            let field_name_str = field_ident.to_string();

            quote! {
                #idx => Some(#field_name_str)
            }
        });

        quote! {
            match idx {
                #(#field_arms,)*
                _ => None,
            }
        }
    } else {
        quote!(None)
    }
}

/// Generates a token stream that implements Reflect and Struct for the provided struct
pub fn derive_reflect_struct(input: &DeriveInput, data_struct: &DataStruct) -> proc_macro2::TokenStream {
    let type_path = &input.ident;
    let name = type_path.to_string();
    //let name = type_path.span().source_text().unwrap();

    let field_len = data_struct.fields.len();
    let get_field_match = gen_struct_field_match(data_struct, false);
    let get_field_mut_match = gen_struct_field_match(data_struct, true);
    let get_field_match_idx = gen_struct_field_match_idx(data_struct, false);
    let get_field_mut_match_idx = gen_struct_field_match_idx(data_struct, true);
    let set_field_named = gen_struct_set_field_match(data_struct);
    let set_field_idx = gen_struct_set_field_match_idx(data_struct);
    let get_field_name_match = gen_struct_field_name_match(data_struct);
    let get_field_name_match_idx = gen_struct_field_name_match_idx(data_struct);
    let field_name_at = gen_struct_field_name_idx(data_struct);

    let generics = add_trait_bounds(input.generics.clone(), vec![parse_quote!(Reflect), parse_quote!(Clone)]);
    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();

    quote! {
        impl #impl_generics lyra_engine::reflect::Reflect for #type_path #ty_generics #where_clause {
            fn name(&self) -> String {
                #name.to_string()
            }
    
            fn type_id(&self) -> std::any::TypeId {
                std::any::TypeId::of::<#type_path #ty_generics>()
            }
    
            fn as_any(&self) -> &dyn std::any::Any {
                self
            }
    
            fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
                self
            }
    
            fn apply(&mut self, val: &dyn lyra_engine::reflect::Reflect) {
                let val = val.as_any().downcast_ref::<Self>()
                    .expect(&format!("`self` was not {}", #name));
                *self = val.clone();
            }
    
            fn clone_inner(&self) -> Box<dyn lyra_engine::reflect::Reflect> {
                Box::new(self.clone())
            }
    
            fn reflect_ref(&self) -> lyra_engine::reflect::ReflectRef {
                lyra_engine::reflect::ReflectRef::Struct(self)
            }
    
            fn reflect_mut(&mut self) -> lyra_engine::reflect::ReflectMut {
                lyra_engine::reflect::ReflectMut::Struct(self)
            }

            fn reflect_val(&self) -> &dyn lyra_engine::reflect::Reflect {
                self
            }

            fn reflect_val_mut(&mut self) -> &mut dyn lyra_engine::reflect::Reflect {
                self
            }
        }

        impl #impl_generics lyra_engine::reflect::Struct for #type_path #ty_generics #where_clause {
            fn field(&self, name: &str) -> Option<&dyn lyra_engine::reflect::Reflect> {
                let name = name.to_lowercase();
                let name = name.as_str();
    
                #get_field_match
            }
    
            fn field_mut(&mut self, name: &str) -> Option<&mut dyn lyra_engine::reflect::Reflect> {
                let name = name.to_lowercase();
                let name = name.as_str();
    
                #get_field_mut_match
            }
    
            fn fields_len(&self) -> usize {
                #field_len
            }
    
            fn field_at(&self, idx: usize) -> Option<&dyn lyra_engine::reflect::Reflect> {
                #get_field_match_idx
            }
    
            fn field_at_mut(&mut self, idx: usize) -> Option<&mut dyn lyra_engine::reflect::Reflect> {
                #get_field_mut_match_idx
            }

            fn field_name_at(&self, idx: usize) -> Option<&str> {
                #field_name_at
            }

            fn set_field(&mut self, name: &str, val: &dyn lyra_engine::reflect::Reflect) -> bool {
                #set_field_named
            }

            fn set_field_at(&mut self, idx: usize, val: &dyn lyra_engine::reflect::Reflect) -> bool {
                #set_field_idx
            }

            fn field_type(&self, name: &str) -> Option<&'static str> {
                #get_field_name_match
            }

            fn field_type_at(&self, idx: usize) -> Option<&'static str> {
                #get_field_name_match_idx
            }
    
            fn method(&self, name: &str) -> Option<&lyra_engine::reflect::Method> {
                unimplemented!()
            }
    
            fn method_at(&self, idx: usize) -> Option<&lyra_engine::reflect::Method> {
                unimplemented!()
            }
    
            fn methods_len(&self) -> usize {
                unimplemented!()
            }
        }
    }
}