matched_enums_macro 1.1.0

extern crate alloc;

use alloc::{vec, vec::Vec};

use hashbrown::HashMap;
use proc_macro2::TokenStream;
use quote::quote;
use syn::{parse2, spanned::Spanned, Ident, ItemEnum, Meta, Type};

use crate::{match_attribute::MatchAttribute, top_level_attribute::TopLevelAttribute};

pub struct VariantMatcherPair(pub Ident, pub MatchAttribute);

pub fn get_match_attrs_per_type(
    enum_type: &ItemEnum,
    attribute: &TopLevelAttribute,
) -> syn::Result<HashMap<Type, Vec<VariantMatcherPair>>> {
    const ATTR_NAME_MATCHER: &str = "matches";

    let mut map = HashMap::from_iter(
        attribute
            .value_types
            .iter()
            .map(|val| (val.clone(), vec![])),
    );

    for variant in &enum_type.variants {
        let match_attrs = variant
            .attrs
            .iter()
            .filter(|attr| attr.path().is_ident(ATTR_NAME_MATCHER));

        for match_attr in match_attrs {
            match &match_attr.meta {
                Meta::List(meta) => {
                    let parsed_attr: MatchAttribute = parse2(meta.tokens.clone())?;
                    let attr_type_bind = parsed_attr.type_bind.clone().unwrap_or(
                        attribute
                            .value_types
                            .first()
                            .ok_or(syn::Error::new(enum_type.span(), "The default `value_types` should not be empty"))?
                            .clone(),
                    );


                    map.get_mut(&attr_type_bind).ok_or(syn::Error::new(
                        match_attr.span(),
                        "Cannot bind match to a type which is not present in the value_types."
                    ))?.push(VariantMatcherPair(variant.ident.clone(), parsed_attr));
                },
                _ => return Err(syn::Error::new(match_attr.span(), format_args!(
                    "`{ATTR_NAME_MATCHER}` should be assigned through as a list (e.g., `#[{ATTR_NAME_MATCHER}(0..42)]`)."
                )))
            }
        }
    }

    Ok(map)
}

pub fn implement_full(
    enum_type: &ItemEnum,
    type_spec: &Type,
    match_pairs: &[VariantMatcherPair],
) -> TokenStream {
    let generics = &enum_type.generics;
    let ident = &enum_type.ident;

    let match_arms = match_pairs
        .iter()
        .clone()
        .map(|VariantMatcherPair(_, attr)| &attr.arm);

    let idents = match_pairs
        .iter()
        .clone()
        .map(|VariantMatcherPair(ident, _)| ident);

    quote! {
        impl<#generics> From<#type_spec> for #ident {
            fn from(value: #type_spec) -> Self {
                match value {
                    #(#match_arms => Self::#idents,)*
                }
            }
        }
    }
}

pub fn implement_partial(
    enum_type: &ItemEnum,
    type_spec: &Type,
    match_pairs: &[VariantMatcherPair],
) -> TokenStream {
    let generics = &enum_type.generics;
    let ident = &enum_type.ident;

    let match_arms = match_pairs
        .iter()
        .clone()
        .map(|VariantMatcherPair(_, attr)| &attr.arm);

    let idents = match_pairs
        .iter()
        .clone()
        .map(|VariantMatcherPair(ident, _)| ident);

    quote! {
        impl<#generics> TryFrom<#type_spec> for #ident {
            type Error = ::matched_enums_types::UnmatchedError;

            fn try_from(value: #type_spec) -> core::result::Result<Self, Self::Error> {
                match value {
                    #(#match_arms => Ok(Self::#idents),)*
                    _ => Err(Self::Error{})
                }
            }
        }
    }
}