rust-utils-macros 0.1.0

use super::{
    ok_or_compile_err,
    some_or_compile_err,
    gen_doc_attrs,
    quote_compile_err
};
use if_chain::if_chain;
use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use syn::{
    parse::Parser,
    parse_quote,
    punctuated::Punctuated,
    spanned::Spanned,
    Attribute,
    Error as ParseError,
    Expr, FnArg,
    GenericArgument,
    Ident, ImplItemFn,
    ItemStruct, Lit,
    Meta, PathArguments,
    Result as ParseResult,
    Signature, Stmt,
    Token, TraitItemFn,
    Type, TypePath,
    Visibility
};
use quote::{quote, ToTokens};

// options for the #[chainable] attribute
#[derive(Clone)]
pub struct ChainableOptions {
    pub collapse_options: bool,
    pub use_into_impl: bool,
    pub doc_attrs: Vec<Attribute>
}

impl ChainableOptions {
    pub fn parse<T: Into<TokenStream2>>(is_field: bool, attr_args: T) -> ParseResult<Self> {
        let attr_args = attr_args.into();
        let mut doc_attrs = Vec::new();

        if attr_args.is_empty() {
            Ok(
                Self {
                    collapse_options: false,
                    use_into_impl: false,
                    doc_attrs: vec![]
                }
            )
        }
        else {
            let collapse_options_option = if is_field {
                "collapse_option"
            }
            else {
                "collapse_options"
            };

            let mut collapse_options = false;
            let mut use_into_impl = false;
            let options = <Punctuated<Meta, Token![,]>>::parse_terminated.parse2(attr_args.clone())?;

            for option in options {
                if let Meta::Path(path) = &option {
                    let name = path.require_ident()?;
                    if name == collapse_options_option {
                        collapse_options = true;
                    }
                    else if name == "use_into_impl" {
                        use_into_impl = true;
                    }
                    else {
                        let msg = format!(
                            "Unknown option \"{name}\"! Valid options are \
                            \"{collapse_options_option}\", \"use_into_impl\" \
                            and \"doc\" (struct fields only)"
                        );
                        return Err(ParseError::new(name.span(), msg));
                    }
                }
                else if let Meta::NameValue(meta_val) = &option {
                    let name = meta_val.path.require_ident()?;

                    if name == "doc" && is_field {
                        if_chain! {
                            if let Expr::Lit(literal) = &meta_val.value;
                            if let Lit::Str(str_val) = &literal.lit;

                            then {
                                doc_attrs = gen_doc_attrs(str_val.value());
                            }
                            else {
                                return Err(ParseError::new(attr_args.span(), "Invalid Input!"));
                            }
                        }
                    }
                    else {
                        let msg = format!(
                            "Unknown option \"{name}\"! Valid options are \
                            \"{collapse_options_option}\", \"use_into_impl\" \
                            and \"doc\" (struct fields only)"
                        );
                        return Err(ParseError::new(name.span(), msg));
                    }
                }
                else {
                    return Err(ParseError::new(attr_args.span(), "Invalid Input!"));
                }
            }

            Ok(Self {
                collapse_options,
                use_into_impl,
                doc_attrs
            })
        }
    }

    pub fn gen_method_from_field(&self, name: &Ident, field_type: &Type, type_vis: &Visibility) -> ImplItemFn {
        let mut in_arg: FnArg = parse_quote!(#name:#field_type);
            let mut set_stmt: Stmt = parse_quote!(self.#name = #name;);

            if self.collapse_options {
                if_chain! {
                    if let Type::Path(TypePath { path, .. }) = &field_type;
                    if let Some(type_segment) = path.segments.last();
                    if type_segment.ident == "Option";
                    if let PathArguments::AngleBracketed(type_args) = &type_segment.arguments;
                    if let GenericArgument::Type(inner_type) = type_args.args.first().unwrap();

                    then {
                        if self.use_into_impl {
                            in_arg = parse_quote!(#name:impl core::convert::Into<#inner_type>);
                            set_stmt = parse_quote!(self.#name = Some(#name.into());)
                        }
                        else {
                            in_arg = parse_quote!(#name:#inner_type);
                            set_stmt = parse_quote!(self.#name = Some(#name);)
                        }
                    }
                }
            }
            else if self.use_into_impl {
                in_arg = parse_quote!(#name:impl core::convert::Into<#field_type>);
                set_stmt = parse_quote!(self.#name = #name.into(););
            }

            let attrs_to_add = &self.doc_attrs;

            parse_quote! {
                #(#attrs_to_add)*
                #[must_use]
                #type_vis fn #name(mut self, #in_arg) -> Self {
                    #set_stmt
                    self
                }
            }
    }
}

pub fn chainable_struct_fields(mut struct_def: ItemStruct) -> TokenStream2 {
    // the generated methods
    let mut methods: Vec<ImplItemFn> = Vec::new();

    for field in &mut struct_def.fields {
        let field_name = some_or_compile_err!(field.ident.as_ref(), "Tuple structs are not supported");
        let mut field_attrs = Vec::new();
        let mut options_tokens = TokenStream2::new();

        // does the caller have any attributes on any of the fields?
        // if yes, process and remove them
        for attr in field.attrs.drain(..) {
            let mut my_attr = false;

            if let Meta::List(list) = &attr.meta {
                let name = ok_or_compile_err!(list.path.require_ident());

                if name == "chainable" {
                    options_tokens = list.tokens.clone();
                    my_attr = true;
                }
            }

            if !my_attr {
                field_attrs.push(attr);
            }
        }

        let options = ok_or_compile_err!(ChainableOptions::parse(true, options_tokens));

        methods.push(
            options.gen_method_from_field(field_name, &field.ty, &struct_def.vis)
        );

        field.attrs = field_attrs;
    }

    let type_name = &struct_def.ident;
    let (impl_generics, ty_generics, where_clause) = &struct_def.generics.split_for_impl();

    quote! {
        #struct_def

        impl #impl_generics #type_name #ty_generics #where_clause {
            #(#methods)*
        }
    }
}

pub fn chainable_trait_method(method: TraitItemFn, attr_args: TokenStream) -> TokenStream2 {
    let TraitItemFn {
        mut attrs,
        mut sig,
        ..
    } = method.clone();

    let options = ok_or_compile_err!(ChainableOptions::parse(false, attr_args));
    let fn_name = sig.ident.clone();
    let fn_name_str = fn_name.to_string();

    if fn_name_str.starts_with("set_") || fn_name_str.starts_with("add_") {
        let doc_attrs = extract_doc_attrs(&mut attrs);

        // set the name of the function
        sig.ident = ok_or_compile_err!(syn::parse_str(&fn_name_str[4..]));

        let args = process_fn_signature(options, &mut sig);

        quote! {
            #method

            #(#doc_attrs)*
            #(#attrs)*
            #[must_use]
            #sig {
                self.#fn_name(#(#args),*);
                self
            }
        }
    }
    else {
        quote_compile_err!("The method must start with \"set_\" or \"add_\"!")
    }
}

pub fn chainable_inh_method(method: ImplItemFn, attr_args: TokenStream) -> TokenStream2 {
    let ImplItemFn {
        mut attrs,
        vis,
        defaultness,
        mut sig,
        ..
    } = method.clone();

    let options = ok_or_compile_err!(ChainableOptions::parse(false, attr_args));
    let fn_name = sig.ident.clone();
    let fn_name_str = fn_name.to_string();

    if fn_name_str.starts_with("set_") || fn_name_str.starts_with("add_") {
        let doc_attrs = extract_doc_attrs(&mut attrs);

        // set the name of the function
        sig.ident = ok_or_compile_err!(syn::parse_str(&fn_name_str[4..]));

        let args = process_fn_signature(options, &mut sig);

        quote! {
            #method

            #(#doc_attrs)*
            #(#attrs)*
            #[must_use]
            #vis #defaultness #sig {
                self.#fn_name(#(#args),*);
                self
            }
        }
    }
    else {
        quote_compile_err!("The method must start with \"set_\" or \"add_\"!")
    }
}

// process the function signature and return the arguments required to
// call the setter/adder method and collapse optionals if the macro
// caller wants it done.
fn process_fn_signature(options: ChainableOptions, sig: &mut Signature) -> Vec<TokenStream2> {
    let mut in_args = Vec::new();
    let mut args = Vec::new();
    in_args.push(FnArg::Receiver(parse_quote! { mut self }));

    for arg in &sig.inputs {
        if let FnArg::Typed(typed_arg) = arg {
            let var_name = &typed_arg.pat;
            let var_type = &typed_arg.ty;

            // let's collapse all the Options to their inner types if the caller wants it
            if options.collapse_options {
                if_chain! {
                    if let Type::Path(TypePath { path, .. }) = &*typed_arg.ty;
                    if let Some(type_segment) = path.segments.last();
                    if type_segment.ident == "Option";
                    if let PathArguments::AngleBracketed(type_args) = &type_segment.arguments;
                    if let GenericArgument::Type(inner_type) = type_args.args.first().unwrap();

                    then {
                        let inner_type = if options.use_into_impl {
                            parse_quote!(impl core::convert::Into<#inner_type>)
                        }
                        else {
                            inner_type.clone()
                        };

                        let var_name = &typed_arg.pat;
                        in_args.push(
                            parse_quote! { #var_name:#inner_type }
                        );

                        args.push(
                            if options.use_into_impl {
                                parse_quote! { Some(#var_name.into()) }
                            }
                            else {
                                parse_quote! { Some(#var_name) }
                            }
                        );
                    }
                    else {
                        if_chain! {
                            if options.use_into_impl;
                            if let Type::Path(TypePath { path, .. }) = &*typed_arg.ty;
                            if let Some(type_segment) = path.segments.last();
                            if type_segment.ident != "Result";

                            then {
                                in_args.push(
                                    parse_quote!(#var_name:impl core::convert::Into<#var_type>)
                                );
                                args.push(
                                    parse_quote!(#var_name.into())
                                );
                            }
                            else {
                                in_args.push(FnArg::Typed(typed_arg.clone()));
                                args.push(typed_arg.pat.to_token_stream());
                            }
                        }
                    }
                }
            }
            else if options.use_into_impl {
                if_chain! {
                    if let Type::Path(TypePath { path, .. }) = &*typed_arg.ty;
                    if let Some(type_segment) = path.segments.last();
                    if type_segment.ident != "Result" && type_segment.ident != "Option";

                    then {
                        in_args.push(
                            parse_quote!(#var_name:impl core::convert::Into<#var_type>)
                        );
                        args.push(
                            parse_quote!(#var_name.into())
                        );
                    }
                    else {
                        in_args.push(FnArg::Typed(typed_arg.clone()));
                        args.push(typed_arg.pat.to_token_stream());
                    }
                }
            }
            else {
                in_args.push(FnArg::Typed(typed_arg.clone()));
                args.push(typed_arg.pat.to_token_stream());
            }
        }
    }

    sig.inputs = in_args.into_iter().collect();
    sig.output = parse_quote! { -> Self };
    args
}

fn extract_doc_attrs(attrs: &mut Vec<Attribute>) -> Vec<Attribute> {
    let mut new_attrs = Vec::new();
    let mut doc_attrs = Vec::new();

    for attr in attrs.drain(..) {
        let mut my_attr = false;
        if let Meta::NameValue(meta_val) = &attr.meta {
            if_chain! {
                if let Ok(name) = meta_val.path.require_ident();
                if name == "doc";
                if let Expr::Lit(expr) = &meta_val.value;
                if let Lit::Str(str_expr) = &expr.lit;

                then {
                    let doc_comment = str_expr.value();
                    for line in doc_comment.lines() {
                        doc_attrs.push(
                            parse_quote! { #[doc = #line] }
                        );
                    }

                    my_attr = true;
                }
            }
        }

        if !my_attr {
            new_attrs.push(attr);
        }
    }

    doc_attrs.push(
        parse_quote! {
            #[doc = ""]
        }
    );

    doc_attrs.push(
        parse_quote! {
            #[doc = "Chainable variant"]
        }
    );

    *attrs = new_attrs;
    doc_attrs
}