subenum 1.2.0

#![allow(clippy::needless_doctest_main)]
#![doc = include_str!("../README.md")]
#![no_std]
extern crate alloc;

mod build;
mod derive;
mod r#enum;
mod param;
mod predicate;
mod visitor;

use alloc::{borrow::ToOwned, collections::BTreeMap, string::ToString, vec::Vec};

use derive::Derive;
use heck::ToSnakeCase;
use proc_macro::TokenStream;
use proc_macro2::Ident;
use quote::quote;
use r#enum::Enum;
use syn::punctuated::Punctuated;
use syn::{parse_macro_input, DeriveInput, Field, Meta, Token, Type};

const SUBENUM: &str = "subenum";
const ERR: &str =
    "subenum must be called with a list of identifiers, like `#[subenum(EnumA, EnumB(derive(Clone)))]`";

fn snake_case(field: &Field) -> Ident {
    let ident = field.ident.as_ref().unwrap_or_else(|| {
        // No ident; the Type must be Path. Use that.
        match &field.ty {
            Type::Path(path) => path.path.get_ident().unwrap(),
            _ => unimplemented!("a"),
        }
    });
    Ident::new(&ident.to_string().to_snake_case(), ident.span())
}

/// Remove our #[subenum(...)] attributes from the input.
fn sanitize_input(input: &mut DeriveInput) {
    let data = match input.data {
        syn::Data::Enum(ref mut data) => data,
        _ => panic!("SubEnum may only be used on enums."),
    };

    for variant in data.variants.iter_mut() {
        // TODO: Switch to Vec::drain_filter once stabilized.
        let mut i = 0;
        while i < variant.attrs.len() {
            if variant.attrs[i].path().is_ident(SUBENUM) {
                variant.attrs.remove(i);
            } else {
                i += 1;
            }
        }
    }
}

fn build_enum_map(
    args: Punctuated<Meta, syn::Token![,]>,
    derives: &[Derive],
) -> BTreeMap<Ident, Enum> {
    args.into_iter()
        .map(|meta| match meta {
            Meta::Path(path) => (path.get_ident().expect(ERR).to_owned(), Vec::new()),
            Meta::List(ml) => (
                ml.path.get_ident().expect(ERR).to_owned(),
                ml.parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
                    .expect(ERR)
                    .into_iter()
                    .map(|meta| syn::parse_quote!(#[#meta]))
                    .collect(),
            ),

            _ => panic!("{}", ERR),
        })
        .map(|(ident, attrs)| {
            (
                ident.clone(),
                Enum::new(ident.clone(), attrs, derives.to_owned()),
            )
        })
        .collect()
}

#[proc_macro_attribute]
pub fn subenum(args: TokenStream, tokens: TokenStream) -> TokenStream {
    let args = parse_macro_input!(args with Punctuated::<Meta, syn::Token![,]>::parse_terminated);
    let mut input = parse_macro_input!(tokens as DeriveInput);
    let data = match input.data {
        syn::Data::Enum(ref mut data) => data,
        _ => panic!("subenum may only be used on enums."),
    };

    let mut derives = Vec::new();
    for attr in &input.attrs {
        if attr.path().is_ident("derive") {
            for meta in attr
                .parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
                .expect("b")
            {
                match meta {
                    Meta::Path(path) => {
                        if path.is_ident("PartialEq") {
                            derives.push(Derive::PartialEq);
                        }
                    }
                    _ => unimplemented!("{:?}", meta),
                }
            }
        }
    }
    let mut enums = build_enum_map(args, &derives);

    let mut self_variant_attrs = alloc::vec![Vec::new(); data.variants.len()];
    for (variant, self_attrs) in data.variants.iter().zip(&mut self_variant_attrs) {
        for attribute in &variant.attrs {
            // Check for "subenum", iterate through the idents.
            if attribute.path().is_ident(SUBENUM) {
                for meta in attribute
                    .parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
                    .expect("b")
                {
                    let mut var = variant.clone();

                    let (ident, attrs) = match meta {
                        Meta::Path(ref path) => (path.get_ident().unwrap().to_owned(), Vec::new()),
                        Meta::List(ml) => (
                            ml.path.get_ident().unwrap().to_owned(),
                            ml.parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
                                .expect(ERR)
                                .into_iter()
                                .map(|meta| syn::parse_quote!(#[#meta]))
                                .collect(),
                        ),
                        _ => unimplemented!("e"),
                    };

                    // We want all attributes except the "subenum" one.
                    var.attrs.retain(|attr| attribute != attr);

                    if ident == input.ident {
                        self_attrs.extend(attrs);
                        continue;
                    }

                    let e = enums
                        .get_mut(&ident)
                        .expect("All enums to be created must be declared at the top-level subenum attribute");
                    e.variants.push(var);
                    e.variants_attributes.push(attrs);
                }
            }
        }
    }
    for (variants, self_attrs) in data.variants.iter_mut().zip(self_variant_attrs) {
        variants.attrs.extend(self_attrs);
    }

    for e in enums.values_mut() {
        e.compute_generics(&input.generics);
    }

    let attrs = input.attrs.clone();
    let enums: Vec<_> = enums
        .into_values()
        .map(|e| e.build(&mut input, &attrs))
        .collect();

    sanitize_input(&mut input);

    quote!(
        #input

        #(#enums)*
    )
    .into()
}