ecs/update_macro_base/src/lib.rs

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

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

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

    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
    }

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

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

    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);
                }
            }
        }
    }

    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::new(
                        #entity_ident,
                        unsafe { #component_name.as_mut() },
                    );
                }
            } else {
                quote! {
                    let #query_ident = Query::new(
                        #entity_ident,
                        unsafe { ( #( #component_names.as_mut(), )* ) },
                    );
                }
            },
            query_ident,
        )
    }

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

    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);
    }
}

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

impl Update {
    pub const MIN_RESOURCE: usize = 0;
    pub 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(),
        }
    }

    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, )* )
            }
        }
    }

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

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

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

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

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

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

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

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

    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 )*
        }
    }

    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,
            )
        }
    }

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

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

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

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

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

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

pub struct ResourceTokens {
    pub resource_types: TokenStream2,
    pub resource_type_impls: TokenStream2,
    pub resources: TokenStream2,
    pub resource_requirement: TokenStream2,
    pub resource_store: TokenStream2,
    pub resource_idents: Vec<Ident>,
}

pub struct QueryTokens<'a> {
    pub filter_types: TokenStream2,
    pub filter_type_impls: TokenStream2,
    pub query_types: TokenStream2,
    pub query_type_impls: TokenStream2,
    pub queries: &'a [Query],
    pub filter_requirements: TokenStream2,
    pub component_requirements: TokenStream2,
    pub query_count: usize,
    pub verify_dedup: Vec<TokenStream2>,
    pub check_entities: Vec<TokenStream2>,
    pub component_stores: Vec<TokenStream2>,
    pub query_structs: Vec<TokenStream2>,
    pub query_idents: Vec<Ident>,
}

impl Update {
    pub fn resourece_tokens(&self) -> ResourceTokens {
        let resource_types = self.resource_types();
        let resource_type_impls = self.resource_type_impls();

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

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

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

        ResourceTokens {
            resource_types,
            resource_type_impls,
            resources,
            resource_requirement,
            resource_store,
            resource_idents,
        }
    }

    pub fn query_tokens(&self) -> QueryTokens<'_> {
        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 queries = &self.queries;

        let filter_requirements = self.filter_requirements();
        let component_requirements = self.component_requirements();

        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);
        }

        QueryTokens {
            filter_types,
            filter_type_impls,
            query_types,
            query_type_impls,
            queries,
            filter_requirements,
            component_requirements,
            query_count,
            verify_dedup,
            check_entities,
            component_stores,
            query_structs,
            query_idents,
        }
    }
}