armour-derive 0.28.0

use proc_macro2::TokenStream;
use quote::quote;
use syn::{DataEnum, Fields, Ident, Variant};

use crate::rapira_field_attrs::extract_idx_attr;

pub fn parse(data: &DataEnum, name: Ident, path: &TokenStream) -> proc_macro::TokenStream {
    let string_name = name.to_string();

    if data.variants.len() > 256 {
        return syn::Error::new_spanned(
            &name,
            "#[derive(GetType)] does not support enums with more than 256 variants",
        )
        .to_compile_error()
        .into();
    }

    let is_simple_enum = data.variants.iter().all(|item| item.fields.is_empty());

    if is_simple_enum {
        return simple_enum(data.variants.iter(), name, path);
    }

    let variants_stream: Vec<TokenStream> = data
        .variants
        .iter()
        .enumerate()
        .map(|(index, variant)| {
            // Wire discriminant: explicit `#[idx = N]` or declaration position — must
            // match rapira-derive (rapira-derive/src/enums.rs), else schema tag diverges
            // from the serialized tag.
            let tag = extract_idx_attr(&variant.attrs)
                .map(|i| i as u8)
                .unwrap_or(index as u8);

            let variant_name = variant.ident.to_string();

            match &variant.fields {
                Fields::Named(fields) => {
                    // Named variant fields are idx-sorted on the wire (rapira-derive),
                    // so the schema must list them in the same order.
                    let mut entries: Vec<(u32, TokenStream)> = Vec::new();
                    let mut seq = 0u32;
                    for field in fields.named.iter() {
                        let typ = &field.ty;
                        let fname = field.ident.as_ref().unwrap().to_string();
                        let fidx = extract_idx_attr(&field.attrs).unwrap_or_else(|| {
                            let current = seq;
                            seq += 1;
                            current
                        });
                        entries.push((
                            fidx,
                            quote! {
                                (#fname, <#typ as #path::GetType>::TYPE),
                            },
                        ));
                    }
                    entries.sort_by_key(|(idx, _)| *idx);
                    let items: Vec<TokenStream> = entries.into_iter().map(|(_, ts)| ts).collect();

                    quote! {
                        (#tag, (#variant_name, #path::Typ::Struct(#path::StructType {
                            name: "",
                            fields: #path::Fields::Named(&[ #(#items)* ]),
                        }))),
                    }
                }
                Fields::Unnamed(fields) => {
                    let fields = &fields.unnamed;

                    if fields.len() == 1 {
                        let field = fields.iter().next().unwrap();
                        let typ = &field.ty;

                        quote! {
                            (#tag, (#variant_name, <#typ as #path::GetType>::TYPE)),
                        }
                    } else {
                        let items = fields
                            .iter()
                            .map(|field| {
                                let typ = &field.ty;

                                quote! {
                                    <#typ as #path::GetType>::TYPE,
                                }
                            })
                            .collect::<Vec<_>>();

                        quote! {
                            (#tag, (#variant_name, #path::Typ::Struct(#path::StructType {
                                name: "",
                                fields: #path::Fields::Unnamed(&[ #(#items)* ]),
                            }))),
                        }
                    }
                }
                Fields::Unit => {
                    quote! {
                        (#tag, (#variant_name, #path::Typ::Scalar(#path::ScalarTyp::Void))),
                    }
                }
            }
        })
        .collect();

    let res = quote! {
        impl #path::GetType for #name {
            const TYPE: #path::Typ = #path::Typ::Enum(#path::EnumType {
                name: #string_name,
                variants: &[
                    #(#variants_stream)*
                ]
            });
        }
    };

    proc_macro::TokenStream::from(res)
}

fn simple_enum<'a>(
    variants: impl Iterator<Item = &'a Variant>,
    name: Ident,
    path: &TokenStream,
) -> proc_macro::TokenStream {
    let string_name = name.to_string();

    let variants_stream: Vec<TokenStream> = variants
        .enumerate()
        .map(|(index, variant)| {
            // Same idx-aware discriminant as the general path.
            let tag = extract_idx_attr(&variant.attrs)
                .map(|i| i as u8)
                .unwrap_or(index as u8);

            let variant_name = variant.ident.to_string();

            quote!((#tag, #variant_name),)
        })
        .collect();

    let res = quote! {
        impl #path::GetType for #name {
            const TYPE: #path::Typ = #path::Typ::SimpleEnum(#path::SimpleEnumType {
                name: #string_name,
                variants: &[
                    #(#variants_stream)*
                ]
            });
        }
    };

    proc_macro::TokenStream::from(res)
}