fayalite/crates/fayalite-proc-macros-impl/src/hdl_enum.rs

// SPDX-License-Identifier: LGPL-3.0-or-later
// See Notices.txt for copyright information
use crate::{
    hdl_type_common::{
        common_derives, get_target, ItemOptions, MakeHdlTypeExpr, MaybeParsed, ParsedGenerics,
        ParsedType, SplitForImpl, TypesParser, WrappedInConst,
    },
    kw, Errors, HdlAttr, PairsIterExt,
};
use proc_macro2::TokenStream;
use quote::{format_ident, quote_spanned, ToTokens};
use syn::{
    parse_quote_spanned,
    punctuated::{Pair, Punctuated},
    token::{Brace, Paren},
    Attribute, Field, FieldMutability, Fields, FieldsNamed, FieldsUnnamed, Generics, Ident,
    ItemEnum, ItemStruct, Token, Type, Variant, Visibility,
};

crate::options! {
    #[options = VariantOptions]
    pub(crate) enum VariantOption {}
}

crate::options! {
    #[options = FieldOptions]
    pub(crate) enum FieldOption {}
}

#[derive(Clone, Debug)]
pub(crate) struct ParsedVariantField {
    pub(crate) paren_token: Paren,
    pub(crate) attrs: Vec<Attribute>,
    pub(crate) options: HdlAttr<FieldOptions, kw::hdl>,
    pub(crate) ty: MaybeParsed<ParsedType, Type>,
    pub(crate) comma_token: Option<Token![,]>,
}

#[derive(Clone, Debug)]
pub(crate) struct ParsedVariant {
    pub(crate) attrs: Vec<Attribute>,
    pub(crate) options: HdlAttr<VariantOptions, kw::hdl>,
    pub(crate) ident: Ident,
    pub(crate) field: Option<ParsedVariantField>,
}

impl ParsedVariant {
    fn parse(
        errors: &mut Errors,
        variant: Variant,
        generics: &MaybeParsed<ParsedGenerics, Generics>,
    ) -> Self {
        let Variant {
            mut attrs,
            ident,
            fields,
            discriminant,
        } = variant;
        let options = errors
            .unwrap_or_default(HdlAttr::parse_and_take_attr(&mut attrs))
            .unwrap_or_default();
        let field = match fields {
            Fields::Unnamed(FieldsUnnamed {
                paren_token,
                unnamed,
            }) if unnamed.len() == 1 => {
                let (field, comma_token) = unnamed.into_pairs().next().unwrap().into_tuple();
                let Field {
                    mut attrs,
                    vis,
                    mutability,
                    ident: _,
                    colon_token: _,
                    ty,
                } = field;
                let options = errors
                    .unwrap_or_default(HdlAttr::parse_and_take_attr(&mut attrs))
                    .unwrap_or_default();
                if !matches!(vis, Visibility::Inherited) {
                    errors.error(
                        &vis,
                        "enum variant fields must not have a visibility specifier",
                    );
                }
                if !matches!(mutability, FieldMutability::None) {
                    // FIXME: use mutability as the spanned tokens,
                    // blocked on https://github.com/dtolnay/syn/issues/1717
                    errors.error(&ty, "field mutability is not supported");
                }
                Some(ParsedVariantField {
                    paren_token,
                    attrs,
                    options,
                    ty: TypesParser::maybe_run(generics.as_ref(), ty, errors),
                    comma_token,
                })
            }
            Fields::Unit => None,
            Fields::Unnamed(fields) if fields.unnamed.is_empty() => None,
            Fields::Named(fields) if fields.named.is_empty() => None,
            Fields::Unnamed(_) | Fields::Named(_) => {
                errors.error(
                    fields,
                    "enum variant must either have no fields or a single parenthesized field",
                );
                None
            }
        };
        if let Some((eq, _)) = discriminant {
            errors.error(eq, "custom enum discriminants are not allowed");
        }
        Self {
            attrs,
            options,
            ident,
            field,
        }
    }
}

#[derive(Clone, Debug)]
pub(crate) struct ParsedEnum {
    pub(crate) attrs: Vec<Attribute>,
    pub(crate) options: HdlAttr<ItemOptions, kw::hdl>,
    pub(crate) vis: Visibility,
    pub(crate) enum_token: Token![enum],
    pub(crate) ident: Ident,
    pub(crate) generics: MaybeParsed<ParsedGenerics, Generics>,
    pub(crate) brace_token: Brace,
    pub(crate) variants: Punctuated<ParsedVariant, Token![,]>,
    pub(crate) match_variant_ident: Ident,
}

impl ParsedEnum {
    fn parse(item: ItemEnum) -> syn::Result<Self> {
        let ItemEnum {
            mut attrs,
            vis,
            enum_token,
            ident,
            mut generics,
            brace_token,
            variants,
        } = item;
        let mut errors = Errors::new();
        let mut options = errors
            .unwrap_or_default(HdlAttr::<ItemOptions, kw::hdl>::parse_and_take_attr(
                &mut attrs,
            ))
            .unwrap_or_default();
        errors.ok(options.body.validate());
        let ItemOptions {
            outline_generated: _,
            target: _,
            custom_bounds,
            no_static: _,
            no_runtime_generics: _,
        } = options.body;
        attrs.retain(|attr| {
            if attr.path().is_ident("repr") {
                errors.error(attr, "#[repr] is not supported on #[hdl] enums");
                false
            } else {
                true
            }
        });
        let generics = if custom_bounds.is_some() {
            MaybeParsed::Unrecognized(generics)
        } else if let Some(generics) = errors.ok(ParsedGenerics::parse(&mut generics)) {
            MaybeParsed::Parsed(generics)
        } else {
            MaybeParsed::Unrecognized(generics)
        };
        let variants = Punctuated::from_iter(
            variants
                .into_pairs()
                .map_pair_value(|v| ParsedVariant::parse(&mut errors, v, &generics)),
        );
        errors.finish()?;
        Ok(Self {
            attrs,
            options,
            vis,
            enum_token,
            generics,
            brace_token,
            variants,
            match_variant_ident: format_ident!("__{}__MatchVariant", ident),
            ident,
        })
    }
}

impl ToTokens for ParsedEnum {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Self {
            attrs,
            options,
            vis,
            enum_token,
            ident,
            generics,
            brace_token,
            variants,
            match_variant_ident,
        } = self;
        let span = ident.span();
        let ItemOptions {
            outline_generated: _,
            target,
            custom_bounds: _,
            no_static,
            no_runtime_generics,
        } = &options.body;
        let target = get_target(target, ident);
        let mut struct_attrs = attrs.clone();
        struct_attrs.push(common_derives(span));
        struct_attrs.push(parse_quote_spanned! {span=>
            #[allow(non_snake_case)]
        });
        let struct_fields = Punctuated::from_iter(variants.pairs().map_pair_value_ref(
            |ParsedVariant {
                 attrs: _,
                 options,
                 ident,
                 field,
             }| {
                let VariantOptions {} = options.body;
                let colon_token;
                let ty = if let Some(ParsedVariantField {
                    paren_token,
                    attrs: _,
                    options,
                    ty,
                    comma_token: _,
                }) = field
                {
                    let FieldOptions {} = options.body;
                    colon_token = Token![:](paren_token.span.open());
                    ty.clone().into()
                } else {
                    colon_token = Token![:](span);
                    parse_quote_spanned! {span=>
                        ()
                    }
                };
                Field {
                    attrs: vec![],
                    vis: vis.clone(),
                    mutability: FieldMutability::None,
                    ident: Some(ident.clone()),
                    colon_token: Some(colon_token),
                    ty,
                }
            },
        ));
        ItemStruct {
            attrs: struct_attrs,
            vis: vis.clone(),
            struct_token: Token![struct](enum_token.span),
            ident: ident.clone(),
            generics: generics.into(),
            fields: if struct_fields.is_empty() {
                Fields::Unit
            } else {
                Fields::Named(FieldsNamed {
                    brace_token: *brace_token,
                    named: struct_fields,
                })
            },
            semi_token: None,
        }
        .to_tokens(tokens);
        let (impl_generics, type_generics, where_clause) = generics.split_for_impl();
        if let (MaybeParsed::Parsed(generics), None) = (generics, no_runtime_generics) {
            generics.make_runtime_generics(tokens, vis, ident, &target, |context| {
                let fields: Vec<_> = variants
                    .iter()
                    .map(|ParsedVariant { ident, field, .. }| {
                        if let Some(ParsedVariantField {
                            ty: MaybeParsed::Parsed(ty),
                            ..
                        }) = field
                        {
                            let expr = ty.make_hdl_type_expr(context);
                            quote_spanned! {span=>
                                #ident: #expr,
                            }
                        } else {
                            quote_spanned! {span=>
                                #ident: (),
                            }
                        }
                    })
                    .collect();
                parse_quote_spanned! {span=>
                    #target {
                        #(#fields)*
                    }
                }
            })
        }
        let mut wrapped_in_const = WrappedInConst::new(tokens, span);
        let tokens = wrapped_in_const.inner();
        {
            let mut wrapped_in_const = WrappedInConst::new(tokens, span);
            let tokens = wrapped_in_const.inner();
            let mut enum_attrs = attrs.clone();
            enum_attrs.push(parse_quote_spanned! {span=>
                #[allow(dead_code)]
            });
            ItemEnum {
                attrs: enum_attrs,
                vis: vis.clone(),
                enum_token: *enum_token,
                ident: ident.clone(),
                generics: generics.into(),
                brace_token: *brace_token,
                variants: Punctuated::from_iter(variants.pairs().map_pair_value_ref(
                    |ParsedVariant {
                         attrs,
                         options: _,
                         ident,
                         field,
                     }| Variant {
                        attrs: attrs.clone(),
                        ident: ident.clone(),
                        fields: match field {
                            Some(ParsedVariantField {
                                paren_token,
                                attrs,
                                options: _,
                                ty,
                                comma_token,
                            }) => Fields::Unnamed(FieldsUnnamed {
                                paren_token: *paren_token,
                                unnamed: Punctuated::from_iter([Pair::new(
                                    Field {
                                        attrs: attrs.clone(),
                                        vis: Visibility::Inherited,
                                        mutability: FieldMutability::None,
                                        ident: None,
                                        colon_token: None,
                                        ty: ty.clone().into(),
                                    },
                                    *comma_token,
                                )]),
                            }),
                            None => Fields::Unit,
                        },
                        discriminant: None,
                    },
                )),
            }
            .to_tokens(tokens);
        }
        let mut enum_attrs = attrs.clone();
        enum_attrs.push(parse_quote_spanned! {span=>
            #[allow(dead_code, non_camel_case_types)]
        });
        ItemEnum {
            attrs: enum_attrs,
            vis: vis.clone(),
            enum_token: *enum_token,
            ident: match_variant_ident.clone(),
            generics: generics.into(),
            brace_token: *brace_token,
            variants: Punctuated::from_iter(variants.pairs().map_pair_value_ref(
                |ParsedVariant {
                     attrs,
                     options: _,
                     ident,
                     field,
                 }| Variant {
                    attrs: attrs.clone(),
                    ident: ident.clone(),
                    fields: match field {
                        Some(ParsedVariantField {
                            paren_token,
                            attrs,
                            options: _,
                            ty,
                            comma_token,
                        }) => Fields::Unnamed(FieldsUnnamed {
                            paren_token: *paren_token,
                            unnamed: Punctuated::from_iter([Pair::new(
                                Field {
                                    attrs: attrs.clone(),
                                    vis: Visibility::Inherited,
                                    mutability: FieldMutability::None,
                                    ident: None,
                                    colon_token: None,
                                    ty: parse_quote_spanned! {span=>
                                        ::fayalite::expr::Expr<#ty>
                                    },
                                },
                                *comma_token,
                            )]),
                        }),
                        None => Fields::Unit,
                    },
                    discriminant: None,
                },
            )),
        }
        .to_tokens(tokens);
        let self_token = Token![self](span);
        for (index, ParsedVariant { ident, field, .. }) in variants.iter().enumerate() {
            if let Some(ParsedVariantField { ty, .. }) = field {
                quote_spanned! {span=>
                    #[automatically_derived]
                    impl #impl_generics #target #type_generics
                    #where_clause
                    {
                        #[allow(non_snake_case, dead_code)]
                        #vis fn #ident<__V: ::fayalite::expr::ToExpr<Type = #ty>>(
                            #self_token,
                            v: __V,
                        ) -> ::fayalite::expr::Expr<Self> {
                            ::fayalite::expr::ToExpr::to_expr(
                                &::fayalite::expr::ops::EnumLiteral::new_by_index(
                                    #self_token,
                                    #index,
                                    ::fayalite::__std::option::Option::Some(
                                        ::fayalite::expr::Expr::canonical(
                                            ::fayalite::expr::ToExpr::to_expr(&v),
                                        ),
                                    ),
                                ),
                            )
                        }
                    }
                }
            } else {
                quote_spanned! {span=>
                    #[automatically_derived]
                    impl #impl_generics #target #type_generics
                    #where_clause
                    {
                        #[allow(non_snake_case, dead_code)]
                        #vis fn #ident(#self_token) -> ::fayalite::expr::Expr<Self> {
                            ::fayalite::expr::ToExpr::to_expr(
                                &::fayalite::expr::ops::EnumLiteral::new_by_index(
                                    #self_token,
                                    #index,
                                    ::fayalite::__std::option::Option::None,
                                ),
                            )
                        }
                    }
                }
            }
            .to_tokens(tokens);
        }
        let variants_token = Ident::new("variants", span);
        let from_canonical_body_fields = Vec::from_iter(variants.iter().enumerate().map(
            |(index, ParsedVariant { ident, field, .. })| {
                let ident_str = ident.to_string();
                let val = if field.is_some() {
                    let missing_value_msg = format!("expected variant {ident} to have a field");
                    quote_spanned! {span=>
                        ::fayalite::ty::Type::from_canonical(ty.expect(#missing_value_msg))
                    }
                } else {
                    quote_spanned! {span=>
                        ::fayalite::__std::assert!(ty.is_none());
                    }
                };
                quote_spanned! {span=>
                    #ident: {
                        let ::fayalite::enum_::EnumVariant {
                            name,
                            ty,
                        } = #variants_token[#index];
                        ::fayalite::__std::assert_eq!(&*name, #ident_str);
                        #val
                    },
                }
            },
        ));
        let variant_access_token = Ident::new("variant_access", span);
        let match_active_scope_match_arms = Vec::from_iter(variants.iter().enumerate().map(
            |(index, ParsedVariant { ident, field, .. })| {
                if field.is_some() {
                    quote_spanned! {span=>
                        #index => #match_variant_ident::#ident(
                            ::fayalite::expr::ToExpr::to_expr(
                                &::fayalite::expr::ops::VariantAccess::new_by_index(
                                    #variant_access_token.base(),
                                    #variant_access_token.variant_index(),
                                ),
                            ),
                        ),
                    }
                } else {
                    quote_spanned! {span=>
                        #index => #match_variant_ident::#ident,
                    }
                }
            },
        ));
        let variants_body_variants = Vec::from_iter(variants.iter().map(
            |ParsedVariant {
                 attrs: _,
                 options,
                 ident,
                 field,
             }| {
                let VariantOptions {} = options.body;
                let ident_str = ident.to_string();
                match field {
                    Some(ParsedVariantField { options, .. }) => {
                        let FieldOptions {} = options.body;
                        quote_spanned! {span=>
                            ::fayalite::enum_::EnumVariant {
                                name: ::fayalite::intern::Intern::intern(#ident_str),
                                ty: ::fayalite::__std::option::Option::Some(
                                    ::fayalite::ty::Type::canonical(&#self_token.#ident),
                                ),
                            },
                        }
                    }
                    None => quote_spanned! {span=>
                        ::fayalite::enum_::EnumVariant {
                            name: ::fayalite::intern::Intern::intern(#ident_str),
                            ty: ::fayalite::__std::option::Option::None,
                        },
                    },
                }
            },
        ));
        let variants_len = variants.len();
        quote_spanned! {span=>
            #[automatically_derived]
            impl #impl_generics ::fayalite::ty::Type for #target #type_generics
            #where_clause
            {
                type BaseType = ::fayalite::enum_::Enum;
                type MaskType = ::fayalite::int::Bool;
                type MatchVariant = #match_variant_ident #type_generics;
                type MatchActiveScope = ::fayalite::module::Scope;
                type MatchVariantAndInactiveScope = ::fayalite::enum_::EnumMatchVariantAndInactiveScope<Self>;
                type MatchVariantsIter = ::fayalite::enum_::EnumMatchVariantsIter<Self>;

                fn match_variants(
                    this: ::fayalite::expr::Expr<Self>,
                    source_location: ::fayalite::source_location::SourceLocation,
                ) -> <Self as ::fayalite::ty::Type>::MatchVariantsIter {
                    ::fayalite::module::enum_match_variants_helper(this, source_location)
                }
                fn mask_type(&#self_token) -> <Self as ::fayalite::ty::Type>::MaskType {
                    ::fayalite::int::Bool
                }
                fn canonical(&#self_token) -> ::fayalite::ty::CanonicalType {
                    ::fayalite::ty::CanonicalType::Enum(::fayalite::enum_::Enum::new(::fayalite::enum_::EnumType::variants(#self_token)))
                }
                #[track_caller]
                #[allow(non_snake_case)]
                fn from_canonical(canonical_type: ::fayalite::ty::CanonicalType) -> Self {
                    let ::fayalite::ty::CanonicalType::Enum(enum_) = canonical_type else {
                        ::fayalite::__std::panic!("expected enum");
                    };
                    let #variants_token = ::fayalite::enum_::EnumType::variants(&enum_);
                    ::fayalite::__std::assert_eq!(#variants_token.len(), #variants_len, "enum has wrong number of variants");
                    Self {
                        #(#from_canonical_body_fields)*
                    }
                }
                fn source_location() -> ::fayalite::source_location::SourceLocation {
                    ::fayalite::source_location::SourceLocation::caller()
                }
            }
            #[automatically_derived]
            impl #impl_generics ::fayalite::enum_::EnumType for #target #type_generics
            #where_clause
            {
                fn match_activate_scope(
                    v: <Self as ::fayalite::ty::Type>::MatchVariantAndInactiveScope,
                ) -> (<Self as ::fayalite::ty::Type>::MatchVariant, <Self as ::fayalite::ty::Type>::MatchActiveScope) {
                    let (#variant_access_token, scope) = v.activate();
                    (
                        match #variant_access_token.variant_index() {
                            #(#match_active_scope_match_arms)*
                            #variants_len.. => ::fayalite::__std::panic!("invalid variant index"),
                        },
                        scope,
                    )
                }
                fn variants(&#self_token) -> ::fayalite::intern::Interned<[::fayalite::enum_::EnumVariant]> {
                    ::fayalite::intern::Intern::intern(&[
                        #(#variants_body_variants)*
                    ][..])
                }
            }
        }
        .to_tokens(tokens);
        if let (None, MaybeParsed::Parsed(generics)) = (no_static, &self.generics) {
            let static_generics = generics.clone().for_static_type();
            let (static_impl_generics, static_type_generics, static_where_clause) =
                static_generics.split_for_impl();
            let static_type_body_variants =
                Vec::from_iter(variants.iter().map(|ParsedVariant { ident, field, .. }| {
                    if let Some(_) = field {
                        quote_spanned! {span=>
                            #ident: ::fayalite::ty::StaticType::TYPE,
                        }
                    } else {
                        quote_spanned! {span=>
                            #ident: (),
                        }
                    }
                }));
            let type_properties = format_ident!("__type_properties", span = span);
            let type_properties_variants =
                Vec::from_iter(variants.iter().map(|ParsedVariant { field, .. }| {
                    let variant = if let Some(ParsedVariantField { ty, .. }) = field {
                        quote_spanned! {span=>
                            ::fayalite::__std::option::Option::Some(
                                <#ty as ::fayalite::ty::StaticType>::TYPE_PROPERTIES,
                            )
                        }
                    } else {
                        quote_spanned! {span=>
                            ::fayalite::__std::option::Option::None
                        }
                    };
                    quote_spanned! {span=>
                        let #type_properties = #type_properties.variant(#variant);
                    }
                }));
            quote_spanned! {span=>
                #[automatically_derived]
                impl #static_impl_generics ::fayalite::ty::StaticType
                for #target #static_type_generics
                #static_where_clause
                {
                    const TYPE: Self = Self {
                        #(#static_type_body_variants)*
                    };
                    const MASK_TYPE: <Self as ::fayalite::ty::Type>::MaskType =
                        ::fayalite::int::Bool;
                    const TYPE_PROPERTIES: ::fayalite::ty::TypeProperties = {
                        let #type_properties = ::fayalite::enum_::EnumTypePropertiesBuilder::new();
                        #(#type_properties_variants)*
                        #type_properties.finish()
                    };
                    const MASK_TYPE_PROPERTIES: ::fayalite::ty::TypeProperties =
                        <::fayalite::int::Bool as ::fayalite::ty::StaticType>::TYPE_PROPERTIES;
                }
            }
            .to_tokens(tokens);
        }
    }
}

pub(crate) fn hdl_enum(item: ItemEnum) -> syn::Result<TokenStream> {
    let item = ParsedEnum::parse(item)?;
    let outline_generated = item.options.body.outline_generated;
    let mut contents = item.to_token_stream();
    if outline_generated.is_some() {
        contents = crate::outline_generated(contents, "hdl-enum-");
    }
    Ok(contents)
}