path: root/ecs-macros/src/lib.rs

use std::path::PathBuf as FsPathBuf;

use proc_macro::TokenStream;
use quote::{format_ident, quote, ToTokens};
use syn::spanned::Spanned;
use syn::{
    parse,
    Attribute,
    GenericParam,
    Generics,
    Ident,
    Item,
    ItemEnum,
    ItemStruct,
    ItemUnion,
    Path,
    Type,
};
use toml::value::{Table as TomlTable, Value as TomlValue};

macro_rules! syn_path {
    ($first_segment: ident $(::$segment: ident)*) => {
        ::syn::Path {
            leading_colon: None,
            segments: ::syn::punctuated::Punctuated::from_iter([
                syn_path_segment!($first_segment),
                $(syn_path_segment!($segment),)*
            ])
        }
    };
}

macro_rules! syn_path_segment {
    ($segment: ident) => {
        ::syn::PathSegment {
            ident: ::proc_macro2::Ident::new(
                stringify!($segment),
                ::proc_macro2::Span::call_site(),
            ),
            arguments: ::syn::PathArguments::None,
        }
    };
}

#[proc_macro_derive(Component, attributes(component))]
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();

    let ecs_path = find_engine_ecs_crate_path().unwrap_or_else(|| syn_path!(ecs));

    let (id_or_ids, get_id) = if !item.generics().params.is_empty() {
        let id_lut_ident =
            format_ident!("{}_ID_LUT", item_ident.to_string().to_uppercase());

        let id_lut = quote! {
            static #id_lut_ident: LazyLock<Mutex<HashMap<TypeId, Uid>>> =
                LazyLock::new(|| Mutex::new(HashMap::new()));
        };

        let generics = item.generics().params.iter().map(|param| match param {
            GenericParam::Type(type_param) => type_param.ident.clone(),
            GenericParam::Lifetime(_) => panic!("Lifetime generics are not supported"),
            GenericParam::Const(_) => panic!("Const generics are not supported"),
        });

        let get_id = quote! {
            *#id_lut_ident
                .try_lock()
                .unwrap()
                .entry(TypeId::of::<(#(#generics,)*)>())
                .or_insert_with(|| Uid::new_unique(UidKind::Component))
        };

        (id_lut, get_id)
    } else {
        let id_lazylock_ident =
            format_ident!("{}_ID", item_ident.to_string().to_uppercase());

        let id_lazylock = quote! {
            static #id_lazylock_ident: LazyLock<Uid> = LazyLock::new(|| {
                Uid::new_unique(UidKind::Component)
            });
        };

        let get_id = quote! {
            *#id_lazylock_ident
        };

        (id_lazylock, get_id)
    };

    let mod_ident = format_ident!(
        "__ecs_priv_component_impl_{}",
        item_ident.to_string().to_lowercase()
    );

    quote! {
        mod #mod_ident {
            use ::std::any::{Any, TypeId};
            use ::std::sync::{LazyLock, Mutex};
            use ::std::collections::HashMap;

            use #ecs_path::component::Component;
            use #ecs_path::system::ComponentRefMut;
            use #ecs_path::system::ComponentRef;
            use #ecs_path::uid::{Uid, Kind as UidKind};
            use #ecs_path::system::Input as SystemInput;
            use #ecs_path::type_name::TypeName;

            use super::*;

            #id_or_ids

            impl #impl_generics Component for #item_ident #type_generics
            #where_clause
            {
                type Component = Self;

                type RefMut<'component> = #ref_mut_type;
                type Ref<'component> = #ref_type;

                fn id() -> Uid
                {
                    #get_id
                }

                fn self_id(&self) -> Uid
                {
                    Self::id()
                }

                fn as_any_mut(&mut self) -> &mut dyn Any
                {
                    self
                }

                fn as_any(&self) -> &dyn Any
                {
                    self
                }
            }

            impl #impl_generics SystemInput for #item_ident #type_generics
            #where_clause
            {
            }

            impl #impl_generics TypeName for #item_ident #type_generics
            #where_clause
            {
                fn type_name(&self) -> &'static str
                {
                    std::any::type_name::<Self>()
                }
            }
        }
    }
    .into()
}

#[proc_macro_derive(Sole, attributes(sole))]
pub fn sole_derive(input: TokenStream) -> TokenStream
{
    let item: TypeItem = parse::<Item>(input).unwrap().try_into().unwrap();

    let item_ident = item.ident();

    let sole_attr = item.attribute::<SoleAttribute>("sole").unwrap_or_default();

    let drop_last = sole_attr.drop_last;

    let (impl_generics, type_generics, where_clause) = item.generics().split_for_impl();

    let ecs_path = find_engine_ecs_crate_path().unwrap_or_else(|| syn_path!(ecs));

    quote! {
        impl #impl_generics #ecs_path::sole::Sole for #item_ident #type_generics
        #where_clause
        {
            fn drop_last(&self) -> bool
            {
                #drop_last
            }

            fn as_any_mut(&mut self) -> &mut dyn std::any::Any
            {
                self
            }

            fn as_any(&self) -> &dyn std::any::Any
            {
                self
            }
        }

        impl #impl_generics #ecs_path::type_name::TypeName for #item_ident #type_generics
        #where_clause
        {
            fn type_name(&self) -> &'static str
            {
                std::any::type_name::<Self>()
            }
        }
    }
    .into()
}

trait FromAttribute: Sized
{
    fn from_attribute(attribute: &Attribute) -> Result<Self, syn::Error>;
}

enum TypeItem
{
    Struct(ItemStruct),
    Enum(ItemEnum),
    Union(ItemUnion),
}

impl TypeItem
{
    fn ident(&self) -> &Ident
    {
        match self {
            Self::Struct(struct_item) => &struct_item.ident,
            Self::Enum(enum_item) => &enum_item.ident,
            Self::Union(union_item) => &union_item.ident,
        }
    }

    fn attribute<Attr: FromAttribute>(&self, attr_ident: &str) -> Option<Attr>
    {
        let item_attrs = match &self {
            Self::Struct(struct_item) => &struct_item.attrs,
            &Self::Enum(enum_item) => &enum_item.attrs,
            &Self::Union(union_item) => &union_item.attrs,
        };

        let mut attr: Option<&Attribute> = None;

        for item_attr in item_attrs {
            if attr.is_some() {
                panic!("Expected only one {} attribute", attr_ident);
            }

            if item_attr.path().get_ident()? == attr_ident {
                attr = Some(item_attr);
            }
        }

        Some(Attr::from_attribute(attr?).unwrap())
    }

    fn generics(&self) -> &Generics
    {
        match self {
            Self::Struct(struct_item) => &struct_item.generics,
            Self::Enum(enum_item) => &enum_item.generics,
            Self::Union(union_item) => &union_item.generics,
        }
    }
}

impl TryFrom<Item> for TypeItem
{
    type Error = syn::Error;

    fn try_from(item: Item) -> Result<Self, Self::Error>
    {
        match item {
            Item::Struct(struct_item) => Ok(Self::Struct(struct_item)),
            Item::Enum(enum_item) => Ok(Self::Enum(enum_item)),
            Item::Union(union_item) => Ok(Self::Union(union_item)),
            _ => Err(syn::Error::new(
                item.span(),
                "Expected a struct, a enum or a union",
            )),
        }
    }
}

impl ToTokens for TypeItem
{
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream)
    {
        match self {
            Self::Struct(struct_item) => struct_item.to_tokens(tokens),
            Self::Enum(enum_item) => enum_item.to_tokens(tokens),
            Self::Union(union_item) => union_item.to_tokens(tokens),
        }
    }
}

#[derive(Debug, Default)]
struct SoleAttribute
{
    drop_last: bool,
}

impl FromAttribute for SoleAttribute
{
    fn from_attribute(attribute: &Attribute) -> Result<Self, syn::Error>
    {
        let mut drop_last = false;

        attribute.parse_nested_meta(|meta| {
            if meta
                .path
                .get_ident()
                .is_some_and(|flag| flag == "drop_last")
            {
                drop_last = true;

                return Ok(());
            }

            Err(meta.error("Unrecognized token"))
        })?;

        Ok(Self { drop_last })
    }
}

fn find_engine_ecs_crate_path() -> Option<Path>
{
    let cargo_manifest_dir = FsPathBuf::from(std::env::var("CARGO_MANIFEST_DIR").ok()?);

    let cargo_crate_name = std::env::var("CARGO_CRATE_NAME").ok()?;
    let cargo_pkg_name = std::env::var("CARGO_PKG_NAME").ok()?;

    if cargo_pkg_name == "ecs" && cargo_crate_name != "ecs" {
        // Macro is used by a ecs crate example/test/benchmark
        return Some(syn_path!(ecs));
    }

    let crate_manifest = std::fs::read_to_string(cargo_manifest_dir.join("Cargo.toml"))
        .ok()?
        .parse::<TomlTable>()
        .expect("Failed to parse crate manifest file");

    let package = match crate_manifest.get("package")? {
        TomlValue::Table(package) => Some(package),
        _ => None,
    }?;

    let package_name = match package.get("name")? {
        TomlValue::String(package_name) => Some(package_name),
        _ => None,
    }?;

    if package_name == "ecs" {
        return Some(syn_path!(crate));
    }

    let crate_dependencies = match crate_manifest.get("dependencies")? {
        TomlValue::Table(dependencies) => Some(dependencies),
        _ => None,
    }?;

    crate_dependencies.iter().find_map(|(crate_dep_name, _)| {
        if crate_dep_name == "engine" {
            return Some(syn_path!(engine::ecs));
        }

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