openconfiguration_derive 1.6.0

use proc_macro::TokenStream;
use quote::quote;

#[proc_macro_derive(Visitable, attributes(visitable))]
/// Macro to derive the `Visitable` trait.
///
/// Expects all fields to either implement `Visitable` or to provide a custom ``visit_with`` attribute.
///
/// Example:
///
/// ```rust
/// use crate::support::Visitable;
///
/// #[derive(Visitable)]
/// struct MyStruct {
///    #[visitable(visit_with = "visit_field")]
///   field: String,
/// }
///
/// fn visit_field(field: &mut String, visitor: &mut dyn Visitor) {}
/// ```
pub fn visitable_derive(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();

    visitable_derive_macro(&ast)
}

fn visitable_derive_macro(ast: &syn::DeriveInput) -> TokenStream {
    let name = &ast.ident;

    match &ast.data {
        syn::Data::Struct(data) => impl_macro_struct(name, data),
        syn::Data::Enum(data) => impl_macro_enum(name, data),
        _ => panic!("Expected input to be a struct or enum"),
    }
}

struct FieldAttributes {
    visit_with: Option<syn::Path>,
}

impl syn::parse::Parse for FieldAttributes {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let terminated = input.parse_terminated(RawFieldAttribute::parse, syn::Token!(,))?;

        let mut field_attributes = FieldAttributes { visit_with: None };

        for entry in terminated {
            match entry.key.to_string().as_str() {
                "visit_with" => {
                    let value = match entry.value {
                        syn::Expr::Lit(syn::ExprLit {
                            lit: syn::Lit::Str(str),
                            ..
                        }) => str.parse()?,
                        _ => {
                            return Err(syn::Error::new(
                                entry.key.span(),
                                format!("visit_with must be a string literal"),
                            ))
                        }
                    };

                    field_attributes.visit_with = Some(value);
                }
                value => {
                    return Err(syn::Error::new(
                        entry.key.span(),
                        format!("Unknown attribute: {}", value),
                    ))
                }
            }
        }

        Ok(field_attributes)
    }
}

struct RawFieldAttribute {
    key: syn::Ident,
    _eq_token: syn::Token![=],
    value: syn::Expr,
}

impl syn::parse::Parse for RawFieldAttribute {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        Ok(RawFieldAttribute {
            key: input.parse()?,
            _eq_token: input.parse()?,
            value: input.parse()?,
        })
    }
}

fn impl_macro_struct(name: &syn::Ident, data: &syn::DataStruct) -> TokenStream {
    let fields = data
        .fields
        .iter()
        .map(|field| {
            let name = field.ident.as_ref().unwrap();

            let visit_with = field.attrs.iter().find_map(|attr| {
                let meta = attr.meta.require_list().ok()?;

                if !meta.path.is_ident("visitable") {
                    return None;
                }

                let attributes = meta
                    .parse_args::<FieldAttributes>()
                    .expect("Invalid attributes");

                attributes.visit_with
            });

            if let Some(visit_with) = visit_with {
                quote! {
                    #visit_with(&mut self.#name, visitor)
                }
            } else {
                quote! {
                    crate::support::Visitable::visit_with(&mut self.#name, visitor)
                }
            }
        })
        .collect::<Vec<_>>();

    let gen = quote! {
        impl crate::support::Visitable for #name {
            fn visit_with(&mut self, visitor: &mut dyn crate::support::Visitor) {
                #(#fields);*
            }
        }
    };
    gen.into()
}

fn impl_macro_enum(name: &syn::Ident, data: &syn::DataEnum) -> TokenStream {
    let variants = data
        .variants
        .iter()
        .map(|variant| {
            let variant_name = &variant.ident;

            quote! {
                #name::#variant_name(data) => crate::support::Visitable::visit_with(data, visitor),
            }
        })
        .collect::<Vec<_>>();

    let gen: quote::__private::TokenStream = quote! {
        impl crate::support::Visitable for #name {
            fn visit_with(&mut self, visitor: &mut dyn crate::support::Visitor) {
                match self {
                    #(#variants)*
                }
            }
        }
    };
    gen.into()
}