oors-macros 0.8.0

/*
 * Copyright (c) 2025 eligamii.
 * Licenced under the MIT licence. See the LICENCE file in the project for full licence information
 */

use proc_macro::TokenStream as TokenStream1;
use proc_macro2::{Ident, TokenStream as TokenStream2};
use quote::{format_ident, quote, ToTokens};
use syn::{parse_macro_input, parse_quote, ExprAssign, Field, Fields, FieldsUnnamed, ItemStruct, Type};
use syn::parse::{Parse, ParseStream, Parser};
use syn::punctuated::Punctuated;
use syn::token::Paren;



pub enum ParentAttr {
    Attr {
        parent_type: Type,
        is_parent_foreign: bool
    },
    Default,
}

impl ParentAttr {
    fn get_parent_type(&self) -> Option<&Type> {
        if let ParentAttr::Attr { parent_type, .. } = self {
            Some(parent_type)
        } else {
            None
        }
    }
}

impl Parse for ParentAttr {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        if input.is_empty() { return Ok(ParentAttr::Default) }

        let assign = input.parse::<ExprAssign>()?;
        let name = syn::parse2::<Ident>(assign.left.into_token_stream())?;
        let parent_type = syn::parse2::<Type>(assign.right.into_token_stream())?;

        let is_parent_foreign = match name.to_string() {
            _ if name == "parent" => false,
            _ if name == "extern_parent" => true,
            str => return Err(syn::Error::new_spanned(name, format!("Expected `parent` or `extern_parent`, got `{str}`").as_str()))
        };

        // Filter types to only real ones
        match parent_type {
            Type::BareFn(_) | Type::ImplTrait(_) | Type::Verbatim(_) | Type::TraitObject(_)
            | Type::Infer(_) | Type::Slice(_) | Type::Paren(_) => {
                return Err(syn::Error::new_spanned(parent_type, "Expected concrete, sized type"))
            },
            _ => ()
        }

        Ok(ParentAttr::Attr {
            parent_type,
            is_parent_foreign
        })

    }
}


pub fn make_accessor(field: &&Field, struct_name: &Ident) -> TokenStream2 {
    let field_name = field.ident.as_ref().unwrap();

    let field_ref = &field_name;
    let field_mut = format_ident!("{}_mut", field_name);

    let ty = &field.ty;

    let doc = field.attrs.iter()
        .filter(|p| p.meta.path().is_ident("doc"))
        .collect::<Vec<_>>();

    let doc_sep = if doc.is_empty() {
        ""
    } else {
        "\n\n---"
    };

    let mut_desc = format!("Get a mutable reference to {}.{doc_sep}", field.ident.as_ref().unwrap());
    let ref_desc = format!("Get a reference to {}.{doc_sep}", field.ident.as_ref().unwrap());

    quote! {
        // SAFETY: See documentation of oors::object

        #[doc = #ref_desc]
        #(#doc)*
        fn #field_ref(&self) -> &#ty {
            unsafe {
                let ptr = self as *const Self as *const #struct_name;
                let field_ptr = std::ptr::addr_of!((*ptr).#field_name);
                &*(field_ptr as *const #ty)
            }
        }

        #[doc = #mut_desc]
        #(#doc)*
        fn #field_mut(&mut self) -> &mut #ty {
            unsafe {
                let ptr = self as *mut Self as *mut #struct_name;
                let field_ptr = std::ptr::addr_of_mut!((*ptr).#field_name);
                &mut *(field_ptr as *mut #ty)
            }
        }
    }
}
pub fn make_accessors_trait(struct_name: &Ident, fields: &Vec<&Field>) -> TokenStream2 {
    let trait_name = format_ident!("__{}Accessors", struct_name);
    let accessors = fields.iter().map(|field| make_accessor(field, struct_name));

    let desc = format!("Accessors for {}. Generated by the `#[object]` attribute. Never implement this trait manually", struct_name);

    quote! {
        // SAFETY: See documentation of oors::object
        #[doc(hidden)]
        #[doc = #desc]
        pub unsafe trait #trait_name  {
            #(#accessors)*
        }
        
        unsafe impl<T> #trait_name for T where T: oors::IsA<#struct_name> {}
    }
}


/// (declaration, definition for ObjectBuilder)
pub fn make_builder_func(struct_name: &Ident, field: &&Field) -> (TokenStream2, TokenStream2) {
    let field_name = field.ident.as_ref().unwrap();
    let field_type = &field.ty;
    
    let with_field = format_ident!("with_{}", field_name);

    let doc = field.attrs.iter()
        .filter(|p| p.meta.path().is_ident("doc"))
        .collect::<Vec<_>>();

    let doc_sep = if doc.is_empty() {
        ""
    } else {
        "\n\n---"
    };

    let desc = format!("Initialize {}.{doc_sep}", &field_name);

    let declaration = quote! {
        #[doc = #desc]
        #(#doc)*
        fn #with_field(self, value: #field_type) -> Self;
    };
    
    let definition = quote! {
        // SAFETY: See documentation of oors::object
        fn #with_field(mut self, value: #field_type) -> Self {
            unsafe {
                let ptr = self._value.as_mut_ptr() as *mut #struct_name;
                let field_ptr = std::ptr::addr_of_mut!((*ptr).#field_name);

                // To know if a field is initialized
                let key = std::mem::offset_of!(#struct_name, #field_name);
                if !self._initialized_fields_size.contains_key(&key) {
                    self._initialized_fields_size.insert(
                        key,
                        oors::_align(core::mem::align_of::<#struct_name>(), core::mem::size_of::<#field_type>())
                    );
                } else {
                    field_ptr.drop_in_place();
                }

                field_ptr.write(value);

                self
            }
        }
    };

    (declaration, definition)
}

pub fn make_builder_trait(struct_name: &Ident, fields: &Vec<&Field>) -> TokenStream2 {
    let trait_name = format_ident!("__{}Builder", struct_name);
    let desc = format!("Builder for {}. Generated by the `#[object]` attribute. Never implement this trait manually", struct_name);

    let funcs = fields.iter().map(|field| make_builder_func(struct_name, field));
    let declarations = funcs.clone().map(|func| func.0);
    let definitions = funcs.map(|func| func.1);
    
    quote! {
        #[doc = #desc]
        pub unsafe trait #trait_name {
            #(#declarations)*
        }
        
        unsafe impl<T> #trait_name for oors::ObjectBuilder<T> where T: oors::IsA<#struct_name>, T: oors::Object {
            #(#definitions)*
        }
    }
}

pub fn make_isa_impls(struct_name: &Ident, parent_type: Option<&Type>) -> TokenStream2 {
    let impl_macro_name = format_ident!("__oors_recursive_impl_{}", struct_name);
    let doc = format!(
        "Recursively `impl` the `IsA` of each parent of `{}` and `IsA<{}>` for `T`.\n# Used internally",
        struct_name,
        struct_name
    );


    if let Some(parent_type) = parent_type {
        let parent_type_name = match parent_type {
            Type::Path(path) => &path.path.segments.last().unwrap().ident,
            _ => unreachable!()
        };

        let parent_impl_macro_name = format_ident!("__oors_recursive_impl_{}", parent_type_name);

        quote! {
            #[doc = #doc]
            #[doc(hidden)]
            #[macro_export]
            macro_rules! #impl_macro_name {
                ($t:ty) => {
                    #parent_impl_macro_name!($t);
                    unsafe impl oors::IsA<#struct_name> for $t {}
                }
            }

            #parent_impl_macro_name!(#struct_name);
        }
    } else {
        quote! {
            #[doc = #doc]
            #[macro_export]
            macro_rules! #impl_macro_name {
                ($t:ty) => {
                    unsafe impl oors::IsA<#struct_name> for $t {}
                }
            }
        }
    }
}

pub fn impl_object(attr: TokenStream1, item: TokenStream1) -> TokenStream1 {
    let attr = parse_macro_input!(attr as ParentAttr);
    
    let item_clone = item.clone();
    let mut struct_item = parse_macro_input!(item_clone as ItemStruct);

    let struct_item_name = struct_item.ident.clone();
    
    // Make the accessors for the struct
    let props = struct_item.fields.iter().filter(|p| {
        p.ident.is_some()
    })
        .collect::<Vec<_>>();

    let props_name = props.iter().map(|p| { p.ident.clone().unwrap() }).collect::<Vec<_>>();
    let props_type = props.iter().map(|p| { p.ty.clone() }).collect::<Vec<_>>();

    let accessors = make_accessors_trait(&struct_item_name, &props);

    // Make the required implementations
    let mut foreign_parent_typeids_impl = quote! {};
    let this_typeids_impl;

    let parent_call = if let Some(parent_type) = attr.get_parent_type() {
        quote! {
            &0 => {
                #parent_type::uninit_selective_drop(uninit_self as *mut #parent_type, _init_offsets, size_of_0);
            },
        }
    } else {
        quote!()
    };

    let selective_uninit_drop = quote! {
        // SAFETY: The _init_offset arg tells us every initialized fields of self. Safe as long as _init_offsets
        // have correct values. Will crash the program trying to call .drop() if not.
        // See line 158
        unsafe fn uninit_selective_drop(uninit_self: *mut Self, _init_offsets: &Vec<usize>, size_of_0: Option<usize>) {
            if size_of_0.is_some_and(|p| p == oors::_align(std::mem::align_of::<Self>(), std::mem::size_of::<Self>())) {
                uninit_self.drop_in_place();
                return;
            }

            unsafe {
                for init_offset in _init_offsets.iter() {
                    match init_offset {
                        #(
                            x if x == &std::mem::offset_of!( #struct_item_name, #props_name ) => {
                                (uninit_self as *mut #props_type).drop_in_place();
                            }
                        )*
                        #parent_call
                        _ => ()
                    }
                }
            }
        }
    };

    match attr {
        ParentAttr::Attr { ref parent_type, ref is_parent_foreign } => {
            if *is_parent_foreign {
                foreign_parent_typeids_impl = quote!(unsafe impl oors::Object for #parent_type {});
            }

            this_typeids_impl = quote ! {
                unsafe impl oors::Object for #struct_item_name {
                    fn type_ids() -> std::collections::HashSet<std::any::TypeId> {
                        let mut ids = #parent_type::type_ids();
                        ids.insert(std::any::TypeId::of::<Self>());
                        ids
                    }

                    #selective_uninit_drop
                }
            }
        }

        ParentAttr::Default => {
            this_typeids_impl = quote!(unsafe impl oors::Object for #struct_item_name {
                #selective_uninit_drop
            });
        }
    }



    
    // Makes the *Builder trait for the struct
    let builder_trait = make_builder_trait(&struct_item_name, &props);
    
    // Make the IsA implementation
    let isa_impl = make_isa_impls(&struct_item_name, attr.get_parent_type());


    // Add the parent struct to Self
    if let Some(typ) = attr.get_parent_type() {
        match struct_item.fields {
            Fields::Named(ref mut fields) => {
                fields.named
                    .insert(
                        0,
                        Field::parse_named.parse2(quote! { base: #typ }).unwrap()
                    );
            },

            Fields::Unnamed(ref mut fields) => {
                fields.unnamed
                    .insert(
                        0,
                        Field::parse_unnamed.parse2(quote! { #typ }).unwrap()
                    );
            }

            _ => struct_item.fields = Fields::Unnamed(FieldsUnnamed {
                paren_token: Paren::default(),
                unnamed: {
                    let mut punct = Punctuated::new();
                    punct.push(parse_quote! { #typ });
                    punct
                }
            })
        }
    }

    quote! {
        #[repr(C)]
        #struct_item
        #foreign_parent_typeids_impl
        #this_typeids_impl
        #accessors
        #isa_impl
        #builder_trait

        unsafe impl IsA<#struct_item_name> for #struct_item_name {}
    }.into()
}