enum_macro_gen 0.2.0

//! This crate provides the `EnumMacroGen` derive proc macro for Rust, which
//! simplifies handling variants of an enum. It generates declartive macros
//! according to a given template.
//!
//! ## Example
//!
//! ```rust ignore
//! use enum_macro_gen::EnumMacroGen;
//!
//! #[derive(EnumMacroGen)]
//! #[enum_macro[handle_test={match: $self.handle_$variant($fields);}]]
//! enum Test {
//!     Foo(Item),
//!     Double(Item, Box<Test>),
//!     Bar,
//! }
//! ```
//!
//! Instead of writing a `match` statement to handle each variant of `Test`, you
//! can use the `handle_test!` a macro generated by `EnumMacroGen`.
//!
//! ```rust
//! // <**GENERATED**>
//! macro_rules! handle_test {
//!     ($self:ident, $test:ident) => {
//!         match $test {
//!             Test::Foo(a_0) => {
//!                 $self.handle_foo(a_0);
//!             }
//!             Test::Double(a_0, a_1) => {
//!                 $self.handle_double(a_0, a_1);
//!             }
//!             Test::Bar => {
//!                 $self.handle_bar();
//!             }
//!         }
//!     };
//! }
//! // </**GENERATED**>
//! ```
//!
//! ## Deriving EnumMacroGen
//!
//! To use `EnumMacroGen`, simply add `#[derive(EnumMacroGen)]` above your enum
//! declaration.
//!
//! You can also specify the format of the generated macro with the `enum_macro`
//! attribute. The attribute value should be a token list containing `$variant`
//! and `$fields`, which will be replaced with the variant name and fields,
//! respectively.

use convert_case::{Case, Casing};
use proc_macro::TokenStream;
use proc_macro2::{Group, Literal, Span, TokenStream as TokenStream2, TokenTree};
use proc_macro_error::{abort, proc_macro_error};
use quote::quote;
use syn::{
    braced, bracketed,
    ext::IdentExt,
    parse::{Parse, ParseStream},
    parse_macro_input, DeriveInput, Ident, Token,
};

#[proc_macro_error]
#[proc_macro_derive(EnumMacroGen, attributes(enum_macro))]
pub fn call_match_enum_refs(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let mut enum_macros = vec![];
    for attr in input.attrs {
        match syn::parse2::<EnumMacro>(attr.tokens.clone()) {
            Ok(v) => {
                enum_macros.extend(v.macros);
            }
            Err(err) => abort! {
                attr.tokens,
                "error parsing: {}", err
            },
        }
    }

    let mut expanded = quote! {};

    for enum_macro in enum_macros.into_iter() {
        let macro_name = &enum_macro.name;

        let item = match &input.data {
            syn::Data::Enum(enum_data) => {
                let enum_name = &input.ident;
                let variants = &enum_data.variants;
                let mut arms = Vec::new();

                for variant in variants {
                    let variant_name = &variant.ident;
                    let variant_name_str = variant_name.to_string();
                    let variant_name_case_str = variant_name_str.to_case(Case::Snake);
                    let fields = &variant.fields;
                    let mut pattern_fields = Vec::new();

                    let arm = match fields {
                        syn::Fields::Named(fields) => {
                            for field in fields.named.iter() {
                                let name = &field.ident;
                                pattern_fields.push(quote! { #name });
                            }
                            let fields_tokens = quote! { #(#pattern_fields),* };
                            let body = interpolate(
                                enum_macro.arm.clone(),
                                variant_name_case_str.as_str(),
                                fields_tokens,
                            );

                            quote! { #enum_name::#variant_name{#(#pattern_fields),*} => { #body } }
                        }
                        syn::Fields::Unnamed(_) => {
                            for (idx, _) in fields.iter().enumerate() {
                                let idx = Ident::new(&format!("a_{}", idx), Span::call_site());
                                pattern_fields.push(quote! { #idx });
                            }
                            let fields_tokens = quote! { #(#pattern_fields),* };
                            let body = interpolate(
                                enum_macro.arm.clone(),
                                variant_name_case_str.as_str(),
                                fields_tokens,
                            );

                            quote! { #enum_name::#variant_name(#(#pattern_fields),*) => { #body }
                            }
                        }
                        syn::Fields::Unit => {
                            let fields_tokens = quote! {};
                            let body = interpolate(
                                enum_macro.arm.clone(),
                                variant_name_case_str.as_str(),
                                fields_tokens,
                            );

                            quote! { #enum_name::#variant_name => { #body } }
                        }
                    };

                    arms.push(arm);
                }

                let attr = if enum_macro.export {
                    quote! { #[macro_export] }
                } else {
                    quote! { }
                };

                quote! {
                    #attr
                    #[macro_export]
                    macro_rules! #macro_name {
                        ($self:expr, $var:expr) => {
                            match $var {
                                #(#arms)*
                            }
                        }
                    }
                }
            }
            // If the input is not an enum, return an error
            _ => quote! {
                compile_error!("#[derive(EnumMacroGen)] can only be used on enums");
            },
        };

        expanded = quote! {
            #expanded
            #item
        };
    }

    expanded.into()
}

#[derive(Debug)]
struct EnumMacro {
    macros: Vec<Macro>,
}

#[derive(Debug)]
struct Macro {
    name: Ident,
    arm: TokenStream2,
    export: bool,
}

impl Parse for EnumMacro {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let content;
        bracketed!(content in input);

        let mut macros = vec![];
        while !content.is_empty() {
            let name: Ident = content.parse()?;
            let mut export = false;
            if content.parse::<Token![:]>().is_ok() {
                let modifier: Ident = content.parse()?;
                if modifier.to_string() == "export" {
                    export = true;
                } else {
                    return Err(syn::Error::new(modifier.span(), "expected `export`"));
                }
            }
            content.parse::<Token![=]>()?;

            let mmacro;
            braced!(mmacro in content);

            let _kind: Ident = mmacro.call(Ident::parse_any)?;
            // TODO: check that kind == "match"
            mmacro.parse::<Token![:]>()?;

            let ts: TokenStream2 = mmacro.parse()?;

            macros.push(Macro { name, arm: ts, export });
        }

        Ok(EnumMacro { macros })
    }
}

/// Interpolates the given `TokenStream2` by replacing occurrences of `$variant`
/// with the given `variant` string and replacing occurrences of `$fields` with
/// the given `fields_replacement` `TokenStream2`.
///
/// # Arguments
///
/// * `tokens` - The `TokenStream2` to interpolate.
/// * `variant` - The string to replace occurrences of `$variant` with, so it
///            can be replaced inside literals. For variant BooFoo, handle_boo_foo.
/// * `# variant` - Replace the two tokens '#variant' with the corresponding
///            string literal, BooFoo -> "boo_foo".
/// * `fields_replacement` - The `TokenStream2` to replace occurrences of
///   `$fields` with.
///
/// # Returns
///
/// The interpolated `TokenStream2`.
fn interpolate(
    tokens: TokenStream2,
    variant: &str,
    fields_replacement: TokenStream2,
) -> TokenStream2 {
    let tokens: Vec<_> = tokens.into_iter().collect();
    let mut out_tokens = Vec::new();
    let mut i = 0;

    while i < tokens.len() {
        let v = 'x: {
            match &tokens[i] {
                TokenTree::Group(g) => {
                    let ng = interpolate(g.stream(), variant, fields_replacement.clone());
                    TokenTree::Group(Group::new(g.delimiter(), ng))
                }
                v @ TokenTree::Punct(p) if p.as_char() == '$' => {
                    if i + 1 < tokens.len() {
                        if let TokenTree::Ident(id) = &tokens[i + 1] {
                            if id == "fields" {
                                i += 1;
                                break 'x TokenTree::Group(Group::new(
                                    proc_macro2::Delimiter::None,
                                    fields_replacement.clone(),
                                ));
                            }
                        }
                    }
                    break 'x v.clone();
                }
                v @ TokenTree::Punct(p) if p.as_char() == '#' => {
                    if i + 1 < tokens.len() {
                        if let TokenTree::Ident(id) = &tokens[i + 1] {
                            if id == "variant" {
                                i += 1;
                                break 'x TokenTree::Literal(Literal::string(&variant));
                            }
                        }
                    }
                    break 'x v.clone();
                }
                v @ TokenTree::Ident(id) => {
                    if i + 1 < tokens.len() {
                        if let TokenTree::Punct(p) = &tokens[i + 1] {
                            if p.as_char() == '$' {
                                if i + 2 < tokens.len() {
                                    if let TokenTree::Ident(id2) = &tokens[i + 2] {
                                        if id2 == "variant" {
                                            i += 2;
                                            break 'x TokenTree::Ident(Ident::new(
                                                &format!("{}{}", id.to_string(), variant),
                                                Span::call_site(),
                                            ));
                                        }
                                    }
                                }
                            }
                        }
                    }
                    break 'x v.clone();
                }
                v => v.clone(),
            }
        };
        out_tokens.push(v);
        i += 1;
    }

    TokenStream2::from_iter(out_tokens)
}