deltapatcher-macros 0.2.2

use proc_macro::TokenStream;
use proc_macro2::{Span, TokenStream as TokenStream2};
use quote::quote;
use syn::{parse_macro_input, Data, DeriveInput};

use darling::{FromAttributes, FromDeriveInput};

/// Field-level config via attributes.
#[derive(FromAttributes)]
#[darling(attributes(deltapatcher))]
struct MacroFieldConfig {
    /// For fields with multiple candidate delta types,
    /// use this to specify the delta type to use.
    delta_ty: syn::Type,
}

/// Struct-level config via attributes.
#[derive(FromDeriveInput)]
#[darling(attributes(deltapatcher))]
struct MacroOverallConfig {
    delta_name: Option<syn::Ident>,
    delta_vis: Option<syn::Visibility>,
    delta_derive: Option<Vec<syn::Ident>>,
}

impl MacroOverallConfig {
    fn delta_name_or_default(&self, ty_name: &syn::Ident) -> syn::Ident {
        self.delta_name
            .clone()
            .unwrap_or_else(|| syn::Ident::new(&format!("{ty_name}Delta"), Span::call_site()))
    }
}

/// Helper to detect if `Diff` is also being derived, to avoid duplicate code generation.
fn has_derive_diff(ast: &DeriveInput) -> bool {
    ast.attrs.iter().any(|attr| {
        if attr.path().is_ident("derive") {
            let mut found_diff = false;
            let _ = attr.parse_nested_meta(|meta| {
                if meta.path.is_ident("Diff") {
                    found_diff = true;
                }
                Ok(())
            });
            found_diff
        } else {
            false
        }
    })
}

/// Derives [`Diff`] and [`Patch`] for a struct or enum, generating a corresponding Delta type.
#[proc_macro_derive(Diff, attributes(deltapatcher))]
pub fn derive_diff(input: TokenStream) -> TokenStream {
    let ast = parse_macro_input!(input as DeriveInput);
    expand_macro(ast, true)
}

/// Derives only [`Patch`] for a struct or enum, generating a corresponding Delta type.
#[proc_macro_derive(Patch, attributes(deltapatcher))]
pub fn derive_patch(input: TokenStream) -> TokenStream {
    let ast = parse_macro_input!(input as DeriveInput);
    
    // If both Diff and Patch are derived, let derive(Diff) handle everything
    // to prevent duplicate Delta structs and implementations.
    if has_derive_diff(&ast) {
        return TokenStream::new();
    }
    
    expand_macro(ast, false)
}

fn expand_macro(ast: DeriveInput, generate_diff: bool) -> TokenStream {
    let top_level_config = MacroOverallConfig::from_derive_input(&ast).unwrap();

    match ast.data {
        Data::Enum(e) => expand_enum(ast.ident, e, top_level_config, generate_diff),
        Data::Struct(s) => expand_struct(ast.vis, ast.ident, s, top_level_config, generate_diff),
        Data::Union(_) => Err(syn::Error::new(
            Span::call_site(),
            "Unions are not supported by deltapatcher derives",
        )),
    }
    .unwrap_or_else(syn::Error::into_compile_error)
    .into()
}

fn expand_enum(
    ident: syn::Ident,
    e: syn::DataEnum,
    config: MacroOverallConfig,
    generate_diff: bool,
) -> syn::Result<TokenStream2> {
    let delta_name = config.delta_name_or_default(&ident);
    let delta_derive = config
        .delta_derive
        .map(|vals| quote! { #[derive(#(#vals),*)] })
        .unwrap_or(quote! {});
    let delta_doc = format!("Autogenerated delta for enum [`{ident}`]");

    struct VariantInfo {
        ident: syn::Ident,
        fields: syn::Fields,
        delta_types: Vec<syn::Type>,
    }

    let mut variants: Vec<VariantInfo> = Vec::with_capacity(e.variants.len());
    for variant in &e.variants {
        let mut delta_types = Vec::with_capacity(variant.fields.len());
        for field in &variant.fields {
            if let Ok(field_config) = MacroFieldConfig::from_attributes(&field.attrs) {
                delta_types.push(field_config.delta_ty);
            } else {
                return Err(syn::Error::new_spanned(
                    field,
                    "Cannot infer delta type. Field must be annotated with #[deltapatcher(delta_ty = SomeDelta)].",
                ));
            }
        }
        variants.push(VariantInfo {
            ident: variant.ident.clone(),
            fields: variant.fields.clone(),
            delta_types,
        });
    }

    // ── Build the delta enum variants ──────────────────────────────────────
    let delta_variant_defs = variants.iter().map(|v| {
        let vname = &v.ident;
        let dtypes = &v.delta_types;
        match &v.fields {
            syn::Fields::Named(n) => {
                let names: Vec<_> = n.named.iter().map(|f| f.ident.clone().unwrap()).collect();
                quote! { #vname { #(#names: #dtypes),* } }
            }
            syn::Fields::Unnamed(_) => quote! { #vname(#(#dtypes),*) },
            syn::Fields::Unit => quote! { #vname },
        }
    });

    // ── diff() match arms (Conditional) ───────────────────────────────────
    let diff_impl = if generate_diff {
        let diff_arms = variants.iter().map(|v| {
            let vname = &v.ident;
            let ftypes = v.fields.iter().map(|f| f.ty.clone()).collect::<Vec<_>>();
            let dtypes = &v.delta_types;
            match &v.fields {
                syn::Fields::Named(n) => {
                    let names: Vec<_> = n.named.iter().map(|f| f.ident.clone().unwrap()).collect();
                    let s_names: Vec<_> = names
                        .iter()
                        .map(|n| syn::Ident::new(&format!("s_{n}"), Span::call_site()))
                        .collect();
                    let i_names: Vec<_> = names
                        .iter()
                        .map(|n| syn::Ident::new(&format!("i_{n}"), Span::call_site()))
                        .collect();
                    quote! {
                        (#ident::#vname { #(#names: #s_names),* }, #ident::#vname { #(#names: #i_names),* }) =>
                            #delta_name::#vname {
                                #(#names: <#ftypes as deltapatcher::Diff<#dtypes>>::diff(#s_names, #i_names)),*
                            }
                    }
                }
                syn::Fields::Unnamed(un) => {
                    let count = un.unnamed.len();
                    let s_ids: Vec<_> = (0..count)
                        .map(|i| syn::Ident::new(&format!("s_{i}"), Span::call_site()))
                        .collect();
                    let i_ids: Vec<_> = (0..count)
                        .map(|i| syn::Ident::new(&format!("i_{i}"), Span::call_site()))
                        .collect();
                    quote! {
                        (#ident::#vname(#(#s_ids),*), #ident::#vname(#(#i_ids),*)) =>
                            #delta_name::#vname(
                                #(<#ftypes as deltapatcher::Diff<#dtypes>>::diff(#s_ids, #i_ids)),*
                            )
                    }
                }
                syn::Fields::Unit => quote! {
                    (#ident::#vname, #ident::#vname) => #delta_name::#vname
                },
            }
        });

        quote! {
            impl deltapatcher::Diff<#delta_name> for #ident {
                fn diff(&self, initial: &Self) -> #delta_name {
                    match (self, initial) {
                        #(#diff_arms,)*
                        _ => #delta_name::Replace(self.clone()),
                    }
                }
            }
        }
    } else {
        quote! {}
    };

    // ── patch() match arms ────────────────────────────────────────────────
    let patch_arms = variants.iter().map(|v| {
        let vname = &v.ident;
        let ftypes = v.fields.iter().map(|f| f.ty.clone()).collect::<Vec<_>>();
        let dtypes = &v.delta_types;
        match &v.fields {
            syn::Fields::Named(n) => {
                let names: Vec<_> = n.named.iter().map(|f| f.ident.clone().unwrap()).collect();
                let d_names: Vec<_> = names
                    .iter()
                    .map(|n| syn::Ident::new(&format!("d_{n}"), Span::call_site()))
                    .collect();
                quote! {
                    (#ident::#vname { #(#names),* }, #delta_name::#vname { #(#names: #d_names),* }) => {
                        #(<#ftypes as deltapatcher::Patch<#dtypes>>::patch(#names, #d_names);)*
                    }
                }
            }
            syn::Fields::Unnamed(un) => {
                let count = un.unnamed.len();
                let s_ids: Vec<_> = (0..count)
                    .map(|i| syn::Ident::new(&format!("s_{i}"), Span::call_site()))
                    .collect();
                let d_ids: Vec<_> = (0..count)
                    .map(|i| syn::Ident::new(&format!("d_{i}"), Span::call_site()))
                    .collect();
                quote! {
                    (#ident::#vname(#(#s_ids),*), #delta_name::#vname(#(#d_ids),*)) => {
                        #(<#ftypes as deltapatcher::Patch<#dtypes>>::patch(#s_ids, #d_ids);)*
                    }
                }
            }
            syn::Fields::Unit => quote! {
                (#ident::#vname, #delta_name::#vname) => {}
            },
        }
    });

    // ── aggregate() match arms ────────────────────────────────────────────
    let aggregate_arms = variants.iter().map(|v| {
        let vname = &v.ident;
        let dtypes = &v.delta_types;
        match &v.fields {
            syn::Fields::Named(n) => {
                let names: Vec<_> = n.named.iter().map(|f| f.ident.clone().unwrap()).collect();
                let d_names: Vec<_> = names
                    .iter()
                    .map(|n| syn::Ident::new(&format!("d_{n}"), Span::call_site()))
                    .collect();
                quote! {
                    (Self::#vname { #(#names),* }, Self::#vname { #(#names: #d_names),* }) => {
                        #(<#dtypes as deltapatcher::delta::Delta>::aggregate(#names, #d_names);)*
                    }
                }
            }
            syn::Fields::Unnamed(un) => {
                let count = un.unnamed.len();
                let s_ids: Vec<_> = (0..count)
                    .map(|i| syn::Ident::new(&format!("s_{i}"), Span::call_site()))
                    .collect();
                let d_ids: Vec<_> = (0..count)
                    .map(|i| syn::Ident::new(&format!("d_{i}"), Span::call_site()))
                    .collect();
                quote! {
                    (Self::#vname(#(#s_ids),*), Self::#vname(#(#d_ids),*)) => {
                        #(<#dtypes as deltapatcher::delta::Delta>::aggregate(#s_ids, #d_ids);)*
                    }
                }
            }
            syn::Fields::Unit => quote! {
                (Self::#vname, Self::#vname) => {}
            },
        }
    });

    // ── aggregate_owned() match arms ──────────────────────────────────────
    let aggregate_owned_arms = variants.iter().map(|v| {
        let vname = &v.ident;
        let dtypes = &v.delta_types;
        match &v.fields {
            syn::Fields::Named(n) => {
                let names: Vec<_> = n.named.iter().map(|f| f.ident.clone().unwrap()).collect();
                let d_names: Vec<_> = names
                    .iter()
                    .map(|n| syn::Ident::new(&format!("d_{n}"), Span::call_site()))
                    .collect();
                quote! {
                    (Self::#vname { #(#names),* }, Self::#vname { #(#names: #d_names),* }) => {
                        #(<#dtypes as deltapatcher::delta::Delta>::aggregate_owned(#names, #d_names);)*
                    }
                }
            }
            syn::Fields::Unnamed(un) => {
                let count = un.unnamed.len();
                let s_ids: Vec<_> = (0..count)
                    .map(|i| syn::Ident::new(&format!("s_{i}"), Span::call_site()))
                    .collect();
                let d_ids: Vec<_> = (0..count)
                    .map(|i| syn::Ident::new(&format!("d_{i}"), Span::call_site()))
                    .collect();
                quote! {
                    (Self::#vname(#(#s_ids),*), Self::#vname(#(#d_ids),*)) => {
                        #(<#dtypes as deltapatcher::delta::Delta>::aggregate_owned(#s_ids, #d_ids);)*
                    }
                }
            }
            syn::Fields::Unit => quote! {
                (Self::#vname, Self::#vname) => {}
            },
        }
    });

    Ok(quote! {
        #[doc = #delta_doc]
        #delta_derive
        #[allow(dead_code)]
        enum #delta_name {
            #(#delta_variant_defs,)*
            /// The enum variant changed entirely; replace with this value.
            Replace(#ident),
        }

        impl deltapatcher::delta::Delta for #delta_name {
            fn aggregate(&mut self, next: &Self) {
                match (self, next) {
                    (slf, Self::Replace(val)) => *slf = Self::Replace(val.clone()),
                    (Self::Replace(current), next) => {
                        <#ident as deltapatcher::Patch<#delta_name>>::patch(current, next);
                    },
                    #(#aggregate_arms,)*
                    _ => unreachable!("delta variants disagree without a Replace — this is a bug"),
                }
            }

            fn aggregate_owned(&mut self, next: Self) {
                match (self, next) {
                    (slf, Self::Replace(val)) => *slf = Self::Replace(val),
                    (Self::Replace(current), next) => {
                        <#ident as deltapatcher::Patch<#delta_name>>::patch(current, &next);
                    },
                    #(#aggregate_owned_arms,)*
                    _ => unreachable!("delta variants disagree without a Replace — this is a bug"),
                }
            }
        }

        impl deltapatcher::Patch<#delta_name> for #ident {
            fn patch(&mut self, delta: &#delta_name) {
                match (self, delta) {
                    #(#patch_arms,)*
                    (slf, #delta_name::Replace(val)) => *slf = val.clone(),
                    _ => unreachable!("mismatched non-Replace delta variant applied to wrong enum variant — this is a bug"),
                }
            }
        }

        #diff_impl
    })
}

fn expand_struct(
    vis: syn::Visibility,
    ident: syn::Ident,
    s: syn::DataStruct,
    config: MacroOverallConfig,
    generate_diff: bool,
) -> syn::Result<TokenStream2> {
    let mut delta_field_type = Vec::with_capacity(s.fields.len());
    for field in &s.fields {
        if let Ok(field_config) = MacroFieldConfig::from_attributes(&field.attrs) {
            delta_field_type.push(field_config.delta_ty);
        } else {
            return Err(syn::Error::new_spanned(
                field,
                "Cannot infer delta type. Field must be annotated with #[deltapatcher(delta_ty = SomeDelta)].",
            ));
        }
    }

    let delta_name = config.delta_name_or_default(&ident);
    let delta_derive = config
        .delta_derive
        .map(|vals| quote! { #[derive(#(#vals),*)] })
        .unwrap_or(quote! {});
    let delta_vis = config.delta_vis.unwrap_or(vis);
    let delta_doc = format!("Autogenerated delta for struct [`{ident}`]");

    Ok(match s.fields {
        syn::Fields::Named(n) => {
            let field_vis = n.named.iter().map(|field| field.vis.clone());
            let field_name = n
                .named
                .iter()
                .map(|f| f.ident.clone().unwrap())
                .collect::<Vec<_>>();
            let field_type = n.named.iter().map(|f| f.ty.clone()).collect::<Vec<_>>();

            let diff_impl = if generate_diff {
                quote! {
                    impl deltapatcher::Diff<#delta_name> for #ident {
                        fn diff(&self, initial: &Self) -> #delta_name {
                            #delta_name {
                                #(#field_name: <#field_type as deltapatcher::Diff<#delta_field_type>>::diff(&self.#field_name, &initial.#field_name)),*
                            }
                        }
                    }
                }
            } else { quote! {} };

            quote! {
                #[doc = #delta_doc]
                #delta_derive
                #delta_vis struct #delta_name {
                    #(#field_vis #field_name: #delta_field_type),*
                }

                impl deltapatcher::delta::Delta for #delta_name {
                    fn aggregate(&mut self, next: &Self) {
                        #(<#delta_field_type as deltapatcher::delta::Delta>::aggregate(&mut self.#field_name, &next.#field_name);)*
                    }

                    fn aggregate_owned(&mut self, next: Self) {
                        #(<#delta_field_type as deltapatcher::delta::Delta>::aggregate_owned(&mut self.#field_name, next.#field_name);)*
                    }
                }

                impl deltapatcher::Patch<#delta_name> for #ident {
                    fn patch(&mut self, delta: &#delta_name) {
                        #(<#field_type as deltapatcher::Patch<#delta_field_type>>::patch(&mut self.#field_name, &delta.#field_name);)*
                    }
                }

                #diff_impl
            }
        }
        syn::Fields::Unnamed(un) => {
            let field_vis = un.unnamed.iter().map(|field| field.vis.clone());
            let field_idx = (0..un.unnamed.len()).collect::<Vec<_>>();
            let field_type = un.unnamed.iter().map(|f| f.ty.clone()).collect::<Vec<_>>();

            let diff_impl = if generate_diff {
                quote! {
                    impl deltapatcher::Diff<#delta_name> for #ident {
                        fn diff(&self, initial: &Self) -> #delta_name {
                            #delta_name (#(<#field_type as deltapatcher::Diff<#delta_field_type>>::diff(&self.#field_idx, &initial.#field_idx)),*)
                        }
                    }
                }
            } else { quote! {} };

            quote! {
                #[doc = #delta_doc]
                #delta_derive
                #delta_vis struct #delta_name(#(#field_vis #delta_field_type),*);

                impl deltapatcher::delta::Delta for #delta_name {
                    fn aggregate(&mut self, next: &Self) {
                        #(<#delta_field_type as deltapatcher::delta::Delta>::aggregate(&mut self.#field_idx, &next.#field_idx);)*
                    }

                    fn aggregate_owned(&mut self, next: Self) {
                        #(<#delta_field_type as deltapatcher::delta::Delta>::aggregate_owned(&mut self.#field_idx, next.#field_idx);)*
                    }
                }

                impl deltapatcher::Patch<#delta_name> for #ident {
                    fn patch(&mut self, delta: &#delta_name) {
                        #(<#field_type as deltapatcher::Patch<#delta_field_type>>::patch(&mut self.#field_idx, &delta.#field_idx);)*
                    }
                }

                #diff_impl
            }
        }
        syn::Fields::Unit => {
            let diff_impl = if generate_diff {
                quote! {
                    impl deltapatcher::Diff<#delta_name> for #ident {
                        fn diff(&self, _initial: &Self) -> #delta_name { #delta_name }
                    }
                }
            } else { quote! {} };

            quote! {
                #[doc = #delta_doc]
                #delta_derive
                #delta_vis struct #delta_name;

                impl deltapatcher::delta::Delta for #delta_name {
                    fn aggregate(&mut self, _next: &Self) {}
                    fn aggregate_owned(&mut self, _next: Self) {}
                }

                impl deltapatcher::Patch<#delta_name> for #ident {
                    fn patch(&mut self, _delta: &#delta_name) {}
                }

                #diff_impl
            }
        }
    })
}