ecs/update_macros/src/update.rs

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::{ToTokens, format_ident, quote};
use syn::Ident;

struct Query {
    components: Vec<Ident>,
    filter: Ident,
}

impl Query {
    pub fn new(components: impl Iterator<Item = Ident>, filter_id: usize) -> Self {
        Self {
            components: components.collect(),
            filter: format_ident!("Filter{filter_id}"),
        }
    }

    pub fn types(&self) -> TokenStream2 {
        let component_list = self
            .components
            .iter()
            .map(|c| quote! { #c })
            .collect::<Vec<_>>();

        let components = if component_list.len() == 1 {
            let component = &component_list[0];
            quote! { #component }
        } else {
            quote! { ( #( #component_list, )* ) }
        };

        components
    }

    pub fn dedup_checks(&self) -> TokenStream2 {
        let components = &self.components;
        let filter = &self.filter;

        quote! {
            #( #filter::verify_dedup::<#components>(); )*
        }
    }

    pub fn check_entity(&self, index: usize) -> TokenStream2 {
        let components = &self.components;
        let filter = &self.filter;

        quote! {
            {
                let mut complies = true;

                #(
                    if !entity.components.contains::<#components>() {
                        complies = false;
                    }
                )*

                if !#filter::check(entity) {
                    complies = false;
                }

                if complies {
                    comply_queries.push(#index);
                }
            }
        }
    }

    pub fn component_store_and_query(&self, index: usize) -> (TokenStream2, TokenStream2, Ident) {
        let components = &self.components;
        let entity_name = format_ident!("entity_{index}");
        let component_names = self
            .components
            .iter()
            .map(|component| format_ident!("q_{index}_{component}"))
            .collect::<Vec<_>>();
        let query_ident = format_ident!("query_{index}");
        let entity_ident = format_ident!("e_{index}");

        (
            quote! {
                let #entity_name = &entities[#index];
                let #entity_ident = #entity_name.as_entity();

                #(
                    #[allow(non_snake_case)]
                    let #component_names = UnsafeComponentStore::from(
                        #entity_name.get_component::<#components>()?
                    );
                )*
            },
            if component_names.len() == 1 {
                let component_name = &component_names[0];

                quote! {
                    let #query_ident = Query {
                        entity: #entity_ident,
                        c: unsafe { #component_name.as_mut() },
                        filter: PhantomData,
                    };
                }
            } else {
                quote! {
                    let #query_ident = Query {
                        entity: #entity_ident,
                        c: unsafe { ( #( #component_names.as_mut(), )* ) },
                        filter: PhantomData,
                    };
                }
            },
            query_ident,
        )
    }

    pub fn components(&self) -> &Vec<Ident> {
        &self.components
    }

    pub fn filter(&self) -> &Ident {
        &self.filter
    }
}

impl ToTokens for Query {
    fn to_tokens(&self, tokens: &mut TokenStream2) {
        let component_list = self
            .components
            .iter()
            .map(|c| quote! { &mut #c })
            .collect::<Vec<_>>();

        let components = if component_list.len() == 1 {
            let component = &component_list[0];
            quote! { #component }
        } else {
            quote! { ( #( #component_list, )* ) }
        };

        let filter = self.filter.clone();

        TokenStream2::from(quote! {
            Query<#components, #filter>
        })
        .to_tokens(tokens);
    }
}

struct Update {
    queries: Vec<Query>,
    resources: Vec<Ident>,
}

impl Update {
    const MIN_RESOURCE: usize = 0;
    const MIN_COMPONENTS: usize = 1;

    pub fn new(queries: impl Iterator<Item = usize>, resources: usize) -> Self {
        Self {
            queries: queries
                .enumerate()
                .map(|(query, component_count)| {
                    Query::new(
                        (0..component_count).map(|component| format_ident!("C{component}Q{query}")),
                        query,
                    )
                })
                .collect(),
            resources: (0..resources)
                .map(|resource| format_ident!("Res{resource}"))
                .collect(),
        }
    }

    pub fn filter_types(&self) -> TokenStream2 {
        let filter_types = self
            .queries
            .iter()
            .map(|query| query.filter())
            .collect::<Vec<_>>();

        if filter_types.len() == 1 {
            let filter = &filter_types[0];

            quote! {
                #filter
            }
        } else {
            quote! {
               ( #( #filter_types, )* )
            }
        }
    }

    pub fn filter_type_impls(&self) -> TokenStream2 {
        let filter_types = self
            .queries
            .iter()
            .map(|query| query.filter())
            .collect::<Vec<_>>();

        quote! {
            #( #filter_types, )*
        }
    }

    pub fn filter_requirements(&self) -> TokenStream2 {
        let filter_requirements = self
            .queries
            .iter()
            .map(|query| query.filter())
            .collect::<Vec<_>>();

        quote! {
           #( #filter_requirements: CheckFilter + 'static, )*
        }
    }

    pub fn query_types(&self) -> TokenStream2 {
        let query_types = self
            .queries
            .iter()
            .map(|query| query.types())
            .collect::<Vec<_>>();

        quote! {
            ( #( #query_types, )* )
        }
    }

    pub fn query_type_impls(&self) -> TokenStream2 {
        let query_types = self
            .queries
            .iter()
            .map(|query| query.components())
            .flatten()
            .collect::<Vec<_>>();

        quote! {
            #( #query_types, )*
        }
    }

    pub fn component_requirements(&self) -> TokenStream2 {
        let component_requirements = self
            .queries
            .iter()
            .map(|query| {
                query.components.iter().map(|component| {
                    quote! {
                        #component: EntityComponent + ComponentDebug,
                    }
                })
            })
            .flatten()
            .collect::<Vec<_>>();

        quote! {
            #( #component_requirements )*
        }
    }

    pub fn resource_store(&self) -> (TokenStream2, Vec<Ident>) {
        let resource_types = &self.resources;
        let resource_idents = self
            .resources
            .iter()
            .map(|resource| format_ident!("res_{resource}"))
            .collect::<Vec<_>>();

        if self.resources.is_empty() {
            (quote! {}, Vec::new())
        } else {
            (
                quote! {
                    #[allow(non_snake_case)]
                    let (
                        #( #resource_idents, )*
                    ): (
                        #( &mut #resource_types, )*
                    ) = world.resources.get_mut()?;
                },
                resource_idents,
            )
        }
    }

    pub fn resource_types(&self) -> TokenStream2 {
        let resource_types = &self.resources;

        quote! {
            ( #( #resource_types, )* )
        }
    }

    pub fn resource_type_impls(&self) -> TokenStream2 {
        let resource_types = &self.resources;

        quote! {
            #( #resource_types, )*
        }
    }

    pub fn reource_requirement(&self) -> TokenStream2 {
        let resource_types = &self.resources;

        quote! {
            #( #resource_types: ResourceTrait, )*
        }
    }
}

impl ToTokens for Update {
    fn to_tokens(&self, tokens: &mut TokenStream2) {
        // generics to specify trait implementations
        let filter_types = self.filter_types();
        let filter_type_impls = self.filter_type_impls();
        let query_types = self.query_types();
        let query_type_impls = self.query_type_impls();
        let resource_types = self.resource_types();
        let resource_type_impls = self.resource_type_impls();

        // panic!("{resource_type_impls}");

        // function parameter
        let queries = &self.queries;
        let resources = if self.resources.is_empty() {
            quote! {}
        } else {
            let resources = &self.resources;

            quote! {
                #( &mut #resources, )*
            }
        };

        // trait requirements
        let filter_requirements = self.filter_requirements();
        let component_requirements = self.component_requirements();
        let reource_requirement = self.reource_requirement();

        let query_count = self.queries.len();

        let verify_dedup = self
            .queries
            .iter()
            .map(|query| query.dedup_checks())
            .collect::<Vec<_>>();

        let check_entities = self
            .queries
            .iter()
            .enumerate()
            .map(|(index, query)| query.check_entity(index))
            .collect::<Vec<_>>();

        let component_stores_and_queries = self
            .queries
            .iter()
            .enumerate()
            .map(|(index, query)| query.component_store_and_query(index))
            .collect::<Vec<_>>();

        let mut component_stores = Vec::new();
        let mut query_structs = Vec::new();
        let mut query_idents = Vec::new();

        for (component_store, query_struct, query_ident) in component_stores_and_queries {
            component_stores.push(component_store);
            query_structs.push(query_struct);
            query_idents.push(query_ident);
        }

        let (resource_store, resource_idents) = self.resource_store();

        TokenStream2::from(quote! {
            impl<Func, #filter_type_impls #query_type_impls #resource_type_impls> CreateArchetype<#query_types, #resource_types, Func, #filter_types> for Archetype
            where
                Func: Fn(&mut Commands, #( #queries, )* #resources ) -> Result<()> + Send + Sync + Clone + 'static,
                #filter_requirements
                #component_requirements
                #reource_requirement
            {
                fn create(f: Func) -> Self {
                    #( #verify_dedup )*

                    Self {
                        queries: #query_count,

                        check_entity: Box::new(move |entity| {
                            let mut comply_queries = Vec::new();

                            #( #check_entities )*

                            if comply_queries.is_empty() {
                                None
                            } else {
                                Some(comply_queries)
                            }
                        }),

                        create_callback: Box::new(move |entities| {
                            #( #component_stores )*

                            let f = f.clone();

                            Ok(Box::new(move |commands, world| {
                                #( #query_structs )*
                                #resource_store

                                f(commands, #( #query_idents, )* #( #resource_idents, )* )
                            }))
                        }),

                        entities: IndexMap::new(),
                    }
                }
            }

            impl<Func, #filter_type_impls #query_type_impls #resource_type_impls> AddUpdates<#query_types, #resource_types, Func, #filter_types> for Updates
            where
                Func: Fn(&mut Commands, #( #queries, )* #resources ) -> Result<()> + Send + Sync + Clone + 'static,
                #filter_requirements
                #component_requirements
                #reource_requirement
            {
                fn add_update(
                    &mut self,
                    priority: u32,
                    func: Func,
                ) -> Result<()> {
                    self.add(priority, Archetype::create(func))
                }
            }

            impl<Func, #filter_type_impls #query_type_impls #resource_type_impls> AddUpdates<#query_types, #resource_types, Func, #filter_types> for WorldBuilder
            where
                Func: Fn(&mut Commands, #( #queries, )* #resources ) -> Result<()> + Send + Sync + Clone + 'static,
                #filter_requirements
                #component_requirements
                #reource_requirement
            {
                fn add_update(
                    &mut self,
                    priority: u32,
                    func: Func,
                ) -> Result<()> {
                    self.updates.add_update(priority, func)
                }
            }
        })
        .to_tokens(tokens);
    }
}

pub fn update(max_components: usize, max_resources: usize) -> TokenStream {
    let mut queries: Vec<Vec<usize>> = Vec::new();

    for c in Update::MIN_COMPONENTS..=max_components {
        queries.push(vec![c]);
    }

    for x in Update::MIN_COMPONENTS..=max_components {
        for y in Update::MIN_COMPONENTS..=max_components {
            queries.push(vec![x, y]);
        }
    }

    for x in Update::MIN_COMPONENTS..=max_components {
        for y in Update::MIN_COMPONENTS..=max_components {
            for z in Update::MIN_COMPONENTS..=max_components {
                queries.push(vec![x, y, z]);
            }
        }
    }

    let mut resources = Vec::new();

    for resource in Update::MIN_RESOURCE..=max_resources {
        resources.push(resource);
    }

    let mut updates = Vec::new();

    for query in queries.iter() {
        for resource_count in resources.iter() {
            updates.push(Update::new(query.iter().cloned(), *resource_count));
        }
    }

    // updates.push(Update::new(vec![1].into_iter(), 0));

    // let q = quote! {
    //     #( #updates )*
    // };

    // panic!("{q}");

    TokenStream::from(quote! {
        #( #updates )*
    })
}