deltapatcher-macros 0.1.1

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

use darling::{FromAttributes, FromDeriveInput};

/// Field-level config via attributes.
#[derive(FromAttributes)]
#[darling(attributes(deltapatcher))]
struct DiffFieldConfig {
    /// 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 DiffOverallConfig {
    delta_name: Option<syn::Ident>,
    delta_vis: Option<syn::Visibility>,
    delta_derive: Option<Vec<syn::Ident>>,
}

impl DiffOverallConfig {
    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()))
    }
}

// TODO doc comment
#[proc_macro_derive(Differentiable, attributes(deltapatcher))]
pub fn derive_differentiable(input: TokenStream) -> TokenStream {
    let ast = parse_macro_input!(input as DeriveInput);
    let top_level_config = DiffOverallConfig::from_derive_input(&ast).unwrap();

    match ast.data {
        Data::Enum(e) => differentiable_enum(ast.ident, e, top_level_config),
        Data::Struct(s) => differentiable_struct(ast.vis, ast.ident, s, top_level_config),
        Data::Union(u) => differentiable_union(ast.ident, u, top_level_config),
    }.unwrap_or_else(syn::Error::into_compile_error).into()
}

fn differentiable_enum(ident: syn::Ident, e: syn::DataEnum, config: DiffOverallConfig) -> 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}`]");

    // For each variant, collect the field-level delta types and metadata.
    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) = DiffFieldConfig::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 ──────────────────────────────────────
    // One variant per source variant (with per-field delta types), plus a
    // `Replace(SourceEnum)` catch-all for cross-variant transitions.
    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 ─────────────────────────────────────────────────
    // Pattern: `(Self::V { fields.. }, Self::V { fields.. })` → per-field delta
    // Fallback: anything else → Replace(self.clone())
    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();
                // bind `self` fields as `s_field` and `initial` fields as `i_field`
                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::delta::Differentiable<#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::delta::Differentiable<#dtypes>>::diff(#s_ids, #i_ids)),*
                        )
                }
            }
            syn::Fields::Unit => quote! {
                (#ident::#vname, #ident::#vname) => #delta_name::#vname
            },
        }
    });

    // ── patch() match arms ────────────────────────────────────────────────
    // Pattern: `(Self::V { .. }, DeltaName::V { .. })` → per-field patch
    // `Replace(val)` → *self = val.clone()
    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::delta::Differentiable<#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::delta::Differentiable<#dtypes>>::patch(#s_ids, #d_ids);)*
                    }
                }
            }
            syn::Fields::Unit => quote! {
                (#ident::#vname, #delta_name::#vname) => {}
            },
        }
    });

    // ── aggregate() match arms ────────────────────────────────────────────
    // Both same per-field variant → aggregate each field delta.
    // Either side is Replace, or variants differ → next wins as Replace.
    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) {
                    // If next is a Replace, we become that Replace (it wins unconditionally).
                    (slf, Self::Replace(val)) => *slf = Self::Replace(val.clone()),
                    // If we are already a Replace, we stay Replace — the patch arm in
                    // `patch()` will overwrite self before applying `next`, so `next`
                    // is expressed relative to the replaced value, not the original one.
                    // We therefore need to apply `next` onto a clone of the replaced value
                    // to bake it into a single Replace.
                    (Self::Replace(current), next) => {
                        current.patch(next);
                    },
                    // Same per-field variant on both sides: aggregate field-by-field.
                    #(#aggregate_arms,)*
                    // Different non-Replace variants: next must carry a full Replace — but
                    // that arm is already handled above, so this is unreachable in
                    // well-formed usage.
                    _ => 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) => {
                        current.patch(&next);
                    },
                    #(#aggregate_owned_arms,)*
                    _ => unreachable!("delta variants disagree without a Replace — this is a bug"),
                }
            }
        }

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

            fn patch(&mut self, delta: &#delta_name) {
                match (self, delta) {
                    #(#patch_arms,)*
                    (slf, #delta_name::Replace(val)) => *slf = val.clone(),
                    // Variant mismatch (same-type variant on self but different delta variant)
                    // means the variants changed — only Replace should arrive in this case;
                    // flag it as a bug if that invariant is violated.
                    _ => unreachable!("mismatched non-Replace delta variant applied to wrong enum variant — this is a bug"),
                }
            }
        }
    })
}

fn differentiable_struct(vis: syn::Visibility, ident: syn::Ident, s: syn::DataStruct, config: DiffOverallConfig) -> syn::Result<TokenStream2> {
    let mut delta_field_type = Vec::with_capacity(s.fields.len());
    for field in &s.fields {
        if let Ok(field_config) = DiffFieldConfig::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);
    // TODO span for delta derive (i.e. highlight the arg when using invalid type)
    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<_>>();

            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::delta::Differentiable<#delta_name> for #ident {
                    fn diff(&self, initial: &Self) -> #delta_name {
                        #delta_name {
                            #(#field_name: <#field_type as deltapatcher::delta::Differentiable<#delta_field_type>>::diff(&self.#field_name, &initial.#field_name)),*
                        }
                    }
                    fn patch(&mut self, delta: &#delta_name) {
                        #(<#field_type as deltapatcher::delta::Differentiable<#delta_field_type>>::patch(&mut self.#field_name, &delta.#field_name);)*
                    }
                }
            }
        },
        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<_>>();

            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::delta::Differentiable<#delta_name> for #ident {
                    fn diff(&self, initial: &Self) -> #delta_name {
                        #delta_name (#(<#field_type as deltapatcher::delta::Differentiable<#delta_field_type>>::diff(&self.#field_idx, &initial.#field_idx)),*)
                    }
                    fn patch(&mut self, delta: &#delta_name) {
                        #(<#field_type as deltapatcher::delta::Differentiable<#delta_field_type>>::patch(&mut self.#field_idx, &delta.#field_idx);)*
                    }
                }
            }
        },
        syn::Fields::Unit => {
            // unit struct, implement no-op

            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::delta::Differentiable<#delta_name> for #ident {
                    fn diff(&self, _initial: &Self) -> #delta_name { #delta_name }
                    fn patch(&mut self, _delta: &#delta_name) {}
                }
            }
        },
    }.into())
}

fn differentiable_union(_ident: syn::Ident, _u: syn::DataUnion, _config: DiffOverallConfig) -> syn::Result<TokenStream2> {
    unimplemented!("unions not yet supported")
}