meta_merge_core 0.1.0

// TODO: move to better location?

use proc_macro2::{Span, TokenStream};
use quote::{ToTokens, quote};
use std::collections::BTreeMap;
use syn::{self, AttrStyle, Attribute, Expr, Meta, Path, Token, token};

/// A small tree to represent nested meta content in a normalizable way
#[derive(Clone)]
enum MetaTree {
    /// `foo`
    Path(Path),
    /// `foo = <expr>`
    NameValue(Path, Expr),
    /// `foo(...)`
    List(Path, Vec<MetaTree>),
}

impl MetaTree {
    pub fn as_path(&self) -> &Path {
        match self {
            MetaTree::Path(p) => p,
            MetaTree::NameValue(p, _) => p,
            MetaTree::List(p, _) => p,
        }
    }
}

fn path_key(p: &Path) -> String {
    // Good-enough structural key for deduping.
    p.to_token_stream().to_string()
}

impl MetaTree {
    fn from_attr(attr: &Attribute) -> syn::Result<MetaTree> {
        match &attr.meta {
            Meta::Path(p) => Ok(MetaTree::Path(p.clone())),
            Meta::NameValue(nv) => Ok(MetaTree::NameValue(nv.path.clone(), nv.value.clone())),
            Meta::List(_) => {
                // It’s a list; use parse_nested_meta to walk its contents.
                let top_path = attr.path().clone();
                let mut items = Vec::new();

                attr.parse_nested_meta(|nested| {
                    let key = nested.path.clone();

                    // foo = <expr>
                    if nested.input.peek(Token![=]) {
                        let _eq: Token![=] = nested.input.parse()?;
                        let expr: Expr = nested.input.parse()?;
                        items.push(MetaTree::NameValue(key, expr));
                        return Ok(());
                    }

                    // foo(...)
                    if nested.input.peek(token::Paren) {
                        let mut inner = Vec::new();
                        nested.parse_nested_meta(|nn| {
                            let inner_key = nn.path.clone();

                            // nested foo = <expr>
                            if nn.input.peek(Token![=]) {
                                let _eq: Token![=] = nn.input.parse()?;
                                let expr: Expr = nn.input.parse()?;
                                inner.push(MetaTree::NameValue(inner_key, expr));
                                return Ok(());
                            }

                            // nested foo(...)
                            if nn.input.peek(token::Paren) {
                                let mut deeper = Vec::new();
                                nn.parse_nested_meta(|nnn| {
                                    let deeper_key = nnn.path.clone();

                                    if nnn.input.peek(Token![=]) {
                                        let _eq: Token![=] = nnn.input.parse()?;
                                        let expr: Expr = nnn.input.parse()?;
                                        deeper.push(MetaTree::NameValue(deeper_key, expr));
                                    } else {
                                        deeper.push(MetaTree::Path(deeper_key));
                                    }
                                    Ok(())
                                })?;
                                inner.push(MetaTree::List(inner_key, deeper));
                                return Ok(());
                            }

                            // nested foo (bare path)
                            inner.push(MetaTree::Path(inner_key));
                            Ok(())
                        })?;
                        items.push(MetaTree::List(key, inner));
                        return Ok(());
                    }

                    // foo (bare path)
                    items.push(MetaTree::Path(key));
                    Ok(())
                })?;

                Ok(MetaTree::List(top_path, items))
            }
        }
    }

    fn many_to_token_stream(items: &[MetaTree]) -> TokenStream {
        let mut out: Vec<TokenStream> = Vec::new();
        for it in items {
            match it {
                MetaTree::Path(p) => out.push(quote!(#p)),
                MetaTree::NameValue(p, expr) => out.push(quote!(#p = #expr)),
                MetaTree::List(p, inner) => {
                    let inner_ts = MetaTree::many_to_token_stream(inner);
                    out.push(quote!(#p( #inner_ts )));
                }
            }
        }
        quote!( #(#out),* )
    }

    fn into_attribute(self, style: AttrStyle) -> Attribute {
        // 1) Build tokens for the inner meta (no #[] yet)
        let meta_ts = match self {
            MetaTree::Path(p) => quote!(#p),
            MetaTree::NameValue(p, expr) => quote!(#p = #expr),
            MetaTree::List(p, inner) => {
                let inner_ts = MetaTree::many_to_token_stream(&inner);
                quote!(#p( #inner_ts ))
            }
        };

        // 2) Parse that into a syn::Meta
        let meta: Meta =
            syn::parse2(meta_ts).expect("failed to parse Meta from synthesized tokens");

        //Construct Attribute directly
        Attribute {
            pound_token: Token![#](Span::call_site()),
            style,
            bracket_token: token::Bracket::default(),
            meta,
        }
    }
}

/// Deep-merge two MetaTrees with the same top-level path.
/// Dest wins on conflicts for NameValue; lists are merged recursively; paths de-duped.
fn merge_meta_trees(dst: &mut MetaTree, src: &MetaTree) {
    match (dst, src) {
        (MetaTree::List(dst_path, dst_items), MetaTree::List(src_path, src_items))
            if path_key(dst_path) == path_key(src_path) =>
        {
            #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
            enum Kind {
                Path,
                NameValue,
                List,
            }

            impl From<&MetaTree> for Kind {
                fn from(m: &MetaTree) -> Self {
                    match m {
                        MetaTree::Path(_) => Kind::Path,
                        MetaTree::NameValue(_, _) => Kind::NameValue,
                        MetaTree::List(_, _) => Kind::List,
                    }
                }
            }

            // Map (key, kind) -> index in dst_items
            let mut index: BTreeMap<(String, Kind), usize> = BTreeMap::new();
            for (ix, meta_tree) in dst_items.iter().enumerate() {
                let (path, kind) = (path_key(meta_tree.as_path()), Kind::from(meta_tree));
                index.entry((path, kind)).or_insert(ix);
            }

            for src_meta_tree in src_items {
                match src_meta_tree {
                    MetaTree::Path(p) => {
                        let k = (path_key(p), Kind::Path);
                        index.entry(k).or_insert_with(|| {
                            dst_items.push(MetaTree::Path(p.clone()));
                            dst_items.len() - 1
                        });
                    }
                    MetaTree::NameValue(p, expr) => {
                        let k = (path_key(p), Kind::NameValue);
                        index.entry(k).or_insert_with(|| {
                            dst_items.push(MetaTree::NameValue(p.clone(), expr.clone()));
                            dst_items.len() - 1
                        }); // else keep dst’s value
                    }
                    MetaTree::List(p, inner) => {
                        let k = (path_key(p), Kind::List);
                        if let Some(&idx) = index.get(&k) {
                            if let MetaTree::List(_, ref mut dst_inner) = dst_items[idx] {
                                // Merge into existing list item
                                let mut merged = MetaTree::List(p.clone(), dst_inner.clone());
                                // TODO: use a queue instead of recursion?
                                merge_meta_trees(
                                    &mut merged,
                                    &MetaTree::List(p.clone(), inner.clone()),
                                );
                                if let MetaTree::List(_, new_inner) = merged {
                                    *dst_inner = new_inner;
                                }
                            }
                        } else {
                            dst_items.push(MetaTree::List(p.clone(), inner.clone()));
                            index.insert(k, dst_items.len() - 1);
                        }
                    }
                }
            }
        }
        _ => { /* different keys/forms: no-op, keep dst */ }
    }
}

pub fn merge_attr(dst: &mut Attribute, src: &Attribute) -> syn::Result<()> {
    let mut dst_tree = MetaTree::from_attr(dst)?;
    let src_tree = MetaTree::from_attr(src)?;
    merge_meta_trees(&mut dst_tree, &src_tree);
    dst.clone_from(&dst_tree.into_attribute(dst.style));
    Ok(())
}

#[allow(dead_code)]
/// Deep-merge `src` attributes into `dst` attributes, de-duping by top-level ident/path.
pub fn merge_attr_vecs_deep(dst: &mut Vec<Attribute>, src: &[Attribute]) -> syn::Result<()> {
    #[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
    struct Key {
        path_key: String,
        is_inner: bool,
    }

    let mut idx: BTreeMap<Key, usize> = BTreeMap::new();
    for (i, a) in dst.iter().enumerate() {
        let key = Key {
            path_key: path_key(a.path()),
            is_inner: matches!(a.style, AttrStyle::Inner(_)),
        };
        idx.entry(key).or_insert(i);
    }

    for s in src {
        let k = Key {
            path_key: path_key(s.path()),
            is_inner: matches!(s.style, AttrStyle::Inner(_)),
        };

        if let Some(&i) = idx.get(&k) {
            // Merge into existing attribute
            let dst_attr = dst[i].clone();
            let mut dst_tree = MetaTree::from_attr(&dst_attr)?;
            let src_tree = MetaTree::from_attr(s)?;
            merge_meta_trees(&mut dst_tree, &src_tree);
            // Rebuild attribute with original style
            dst[i] = dst_tree.into_attribute(dst_attr.style);
        } else {
            // No existing: take it verbatim
            dst.push(s.clone());
            idx.insert(k, dst.len() - 1);
        }
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::extensions::token_stream::TokenStreamExt;
    use quote::quote;
    use syn::parse_quote;

    fn pretty_print_or_panic(ts: &TokenStream, message: &str) -> String {
        match ts.clone().to_pretty_print_string() {
            Ok(output) => output,
            Err(err) => {
                panic!("failed to pretty print {message}:\n{ts}\n\n{err}");
            }
        }
    }

    #[inline]
    fn assert_eq_attrs(left: &Vec<Attribute>, right: &Vec<Attribute>) {
        // pretty print requires a valid rust file, so we add a dummy struct to attach attributes to
        let dummy = quote! { struct __DummyStruct__; };
        let left = quote! {
            #(#left)*
            #dummy
        };
        let right = quote! {
            #(#right)*
            #dummy
        };

        let left = pretty_print_or_panic(&left, "left");
        let right = match right.clone().to_pretty_print_string() {
            Ok(output) => output,
            Err(err) => {
                panic!("failed to pretty print right:\n{right}\n\n{err}");
            }
        };

        // remove dummy for comparison
        let dummy = dummy.to_pretty_print_string().unwrap();
        let left = left.replace(&dummy, "");
        let right = right.replace(&dummy, "");

        let left = left.trim();
        let right = right.trim();

        assert_eq!(
            left, right,
            "tokens don't match\n\n/* left */\n\n{left}\n/* right */\n\n{right}\n/* end */"
        );
    }

    #[test]
    fn simple_merge() {
        let src_attrs = vec![parse_quote!( #[foo(bar)] )];
        let mut dst_attrs = vec![parse_quote!( #[foo(bar)] )];
        assert_eq!(
            merge_attr_vecs_deep(&mut dst_attrs, &src_attrs)
                .as_ref()
                .map_err(syn::Error::to_string),
            Ok(&())
        );
        assert_eq_attrs(&dst_attrs, &src_attrs);

        let src_attrs = vec![parse_quote!( #[foo(fizz)] )];
        let mut dst_attrs = vec![parse_quote!( #[foo(bar)] )];
        assert_eq!(
            merge_attr_vecs_deep(&mut dst_attrs, &src_attrs)
                .as_ref()
                .map_err(syn::Error::to_string),
            Ok(&())
        );
        assert_eq_attrs(&dst_attrs, &vec![parse_quote!( #[foo(bar, fizz)] )]);
    }

    #[test]
    fn simple_merge_2_params() {
        let src_attrs = vec![parse_quote!( #[foo(bar, fizz)] )];
        let mut dst_attrs = vec![parse_quote!( #[foo(bar)] )];
        assert_eq!(
            merge_attr_vecs_deep(&mut dst_attrs, &src_attrs)
                .as_ref()
                .map_err(syn::Error::to_string),
            Ok(&())
        );
        assert_eq_attrs(&dst_attrs, &src_attrs);
    }

    #[test]
    fn simple_merge_2_attrs() {
        let src_attrs = vec![
            parse_quote!( #[foo(fizz)] ),
            parse_quote!( #[derive(PartialEq)] ),
        ];
        let mut dst_attrs = vec![
            parse_quote!( #[foo(bar)] ),
            parse_quote!( #[derive(Debug)] ),
        ];
        assert_eq!(
            merge_attr_vecs_deep(&mut dst_attrs, &src_attrs)
                .as_ref()
                .map_err(syn::Error::to_string),
            Ok(&())
        );
        assert_eq_attrs(
            &dst_attrs,
            &vec![
                parse_quote!( #[foo(bar, fizz)] ),
                parse_quote!( #[derive(Debug, PartialEq)] ),
            ],
        );
    }

    #[test]
    fn test_merge_attr() {
        let src_attr = parse_quote!( #[foo(bar)] );
        let mut dst_attr = parse_quote!( #[foo(bar)] );
        assert_eq!(
            merge_attr(&mut dst_attr, &src_attr)
                .as_ref()
                .map_err(syn::Error::to_string),
            Ok(&())
        );
        assert_eq_attrs(&vec![dst_attr], &vec![src_attr]);
    }
}