fayalite/crates/fayalite-proc-macros-impl/src/module/transform_body/expand_match.rs

// SPDX-License-Identifier: LGPL-3.0-or-later
// See Notices.txt for copyright information
use crate::{
    fold::{impl_fold, DoFold},
    module::transform_body::{with_debug_clone_and_fold, Visitor},
    Errors, HdlAttr, PairsIterExt,
};
use proc_macro2::{Span, TokenStream};
use quote::{format_ident, quote_spanned, ToTokens, TokenStreamExt};
use syn::{
    fold::{fold_arm, fold_expr_match, fold_pat, Fold},
    parse::Nothing,
    parse_quote_spanned,
    punctuated::Punctuated,
    spanned::Spanned,
    token::{Brace, Paren},
    Arm, Attribute, Expr, ExprMatch, FieldPat, Ident, Member, Pat, PatIdent, PatOr, PatParen,
    PatPath, PatRest, PatStruct, PatTupleStruct, PatWild, Path, PathSegment, Token, TypePath,
};

with_debug_clone_and_fold! {
    struct MatchPatBinding<> {
        ident: Ident,
    }
}

impl ToTokens for MatchPatBinding {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Self { ident } = self;
        ident.to_tokens(tokens);
    }
}

with_debug_clone_and_fold! {
    struct MatchPatParen<P> {
        paren_token: Paren,
        pat: Box<P>,
    }
}

impl<P: ToTokens> ToTokens for MatchPatParen<P> {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Self { paren_token, pat } = self;
        paren_token.surround(tokens, |tokens| pat.to_tokens(tokens));
    }
}

with_debug_clone_and_fold! {
    struct MatchPatOr<P> {
        leading_vert: Option<Token![|]>,
        cases: Punctuated<P, Token![|]>,
    }
}

impl<P: ToTokens> ToTokens for MatchPatOr<P> {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Self {
            leading_vert,
            cases,
        } = self;
        leading_vert.to_tokens(tokens);
        cases.to_tokens(tokens);
    }
}

with_debug_clone_and_fold! {
    struct MatchPatWild<> {
        underscore_token: Token![_],
    }
}

impl ToTokens for MatchPatWild {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Self { underscore_token } = self;
        underscore_token.to_tokens(tokens);
    }
}

with_debug_clone_and_fold! {
    struct MatchPatStructField<> {
        field_name: Ident,
        colon_token: Option<Token![:]>,
        pat: MatchPatSimple,
    }
}

impl ToTokens for MatchPatStructField {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Self {
            field_name,
            colon_token,
            pat,
        } = self;
        field_name.to_tokens(tokens);
        if let (None, MatchPatSimple::Binding(MatchPatBinding { ident })) = (colon_token, pat) {
            if field_name == ident {
                return;
            }
        }
        colon_token
            .unwrap_or_else(|| Token![:](field_name.span()))
            .to_tokens(tokens);
        pat.to_tokens(tokens);
    }
}

impl MatchPatStructField {
    fn parse(state: &mut HdlMatchParseState<'_>, field_pat: FieldPat) -> Result<Self, ()> {
        let FieldPat {
            attrs: _,
            member,
            colon_token,
            pat,
        } = field_pat;
        let field_name = if let Member::Named(field_name) = member {
            field_name
        } else {
            state
                .errors
                .error(&member, "field name must not be a number");
            format_ident!("_{}", member)
        };
        Ok(Self {
            field_name,
            colon_token,
            pat: MatchPatSimple::parse(state, *pat)?,
        })
    }
}

with_debug_clone_and_fold! {
    struct MatchPatStruct<> {
        match_span: Span,
        path: Path,
        brace_token: Brace,
        fields: Punctuated<MatchPatStructField, Token![,]>,
        rest: Option<Token![..]>,
    }
}

impl ToTokens for MatchPatStruct {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Self {
            match_span,
            path,
            brace_token,
            fields,
            rest,
        } = self;
        quote_spanned! {*match_span=>
            __MatchTy::<#path>
        }
        .to_tokens(tokens);
        brace_token.surround(tokens, |tokens| {
            fields.to_tokens(tokens);
            rest.to_tokens(tokens);
        })
    }
}

with_debug_clone_and_fold! {
    struct MatchPatEnumVariant<> {
        match_span: Span,
        variant_path: Path,
        enum_path: Path,
        variant_name: Ident,
        field: Option<(Paren, MatchPatSimple)>,
    }
}

impl ToTokens for MatchPatEnumVariant {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Self {
            match_span,
            variant_path: _,
            enum_path,
            variant_name,
            field,
        } = self;
        quote_spanned! {*match_span=>
            __MatchTy::<#enum_path>::#variant_name
        }
        .to_tokens(tokens);
        if let Some((paren_token, field)) = field {
            paren_token.surround(tokens, |tokens| field.to_tokens(tokens));
        }
    }
}

#[derive(Debug, Clone)]
enum MatchPatSimple {
    Paren(MatchPatParen<MatchPatSimple>),
    Or(MatchPatOr<MatchPatSimple>),
    Binding(MatchPatBinding),
    Wild(MatchPatWild),
}

impl_fold! {
    enum MatchPatSimple<> {
        Paren(MatchPatParen<MatchPatSimple>),
        Or(MatchPatOr<MatchPatSimple>),
        Binding(MatchPatBinding),
        Wild(MatchPatWild),
    }
}

impl ToTokens for MatchPatSimple {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        match self {
            Self::Or(v) => v.to_tokens(tokens),
            Self::Paren(v) => v.to_tokens(tokens),
            Self::Binding(v) => v.to_tokens(tokens),
            Self::Wild(v) => v.to_tokens(tokens),
        }
    }
}

struct EnumPath {
    variant_path: Path,
    enum_path: Path,
    variant_name: Ident,
}

fn parse_enum_path(variant_path: TypePath) -> Result<EnumPath, TypePath> {
    let TypePath {
        qself: None,
        path: variant_path,
    } = variant_path
    else {
        return Err(variant_path);
    };
    if variant_path.is_ident("HdlNone") || variant_path.is_ident("HdlSome") {
        let ident = variant_path.get_ident().unwrap();
        return Ok(EnumPath {
            enum_path: parse_quote_spanned! {ident.span()=>
                ::fayalite::enum_::HdlOption::<_>
            },
            variant_name: ident.clone(),
            variant_path,
        });
    }
    if variant_path.segments.len() < 2 {
        return Err(TypePath {
            qself: None,
            path: variant_path,
        });
    }
    let mut enum_path = variant_path.clone();
    let PathSegment {
        ident: variant_name,
        arguments,
    } = enum_path.segments.pop().unwrap().into_value();
    if !arguments.is_none() {
        return Err(TypePath {
            qself: None,
            path: variant_path,
        });
    }
    enum_path.segments.pop_punct();
    Ok(EnumPath {
        variant_path,
        enum_path,
        variant_name,
    })
}

fn parse_enum_ident(ident: Ident) -> Result<EnumPath, Ident> {
    parse_enum_path(TypePath {
        qself: None,
        path: ident.into(),
    })
    .map_err(|p| p.path.segments.into_iter().next().unwrap().ident)
}

trait ParseMatchPat: Sized {
    fn simple(v: MatchPatSimple) -> Self;
    fn or(v: MatchPatOr<Self>) -> Self;
    fn paren(v: MatchPatParen<Self>) -> Self;
    fn struct_(state: &mut HdlMatchParseState<'_>, v: MatchPatStruct) -> Result<Self, ()>;
    fn enum_variant(state: &mut HdlMatchParseState<'_>, v: MatchPatEnumVariant)
        -> Result<Self, ()>;
    fn parse(state: &mut HdlMatchParseState<'_>, pat: Pat) -> Result<Self, ()> {
        match pat {
            Pat::Ident(PatIdent {
                attrs: _,
                by_ref,
                mutability,
                ident,
                subpat,
            }) => {
                if let Some(by_ref) = by_ref {
                    state
                        .errors
                        .error(by_ref, "ref not allowed in #[hdl] patterns");
                }
                if let Some(mutability) = mutability {
                    state
                        .errors
                        .error(mutability, "mut not allowed in #[hdl] patterns");
                }
                if let Some((at_token, _)) = subpat {
                    state
                        .errors
                        .error(at_token, "@ not allowed in #[hdl] patterns");
                }
                match parse_enum_ident(ident) {
                    Ok(EnumPath {
                        variant_path,
                        enum_path,
                        variant_name,
                    }) => Self::enum_variant(
                        state,
                        MatchPatEnumVariant {
                            match_span: state.match_span,
                            variant_path,
                            enum_path,
                            variant_name,
                            field: None,
                        },
                    ),
                    Err(ident) => Ok(Self::simple(MatchPatSimple::Binding(MatchPatBinding {
                        ident,
                    }))),
                }
            }
            Pat::Or(PatOr {
                attrs: _,
                leading_vert,
                cases,
            }) => Ok(Self::or(MatchPatOr {
                leading_vert,
                cases: cases
                    .into_pairs()
                    .filter_map_pair_value(|pat| Self::parse(state, pat).ok())
                    .collect(),
            })),
            Pat::Paren(PatParen {
                attrs: _,
                paren_token,
                pat,
            }) => Ok(Self::paren(MatchPatParen {
                paren_token,
                pat: Box::new(Self::parse(state, *pat)?),
            })),
            Pat::Path(PatPath {
                attrs: _,
                qself,
                path,
            }) => {
                let EnumPath {
                    variant_path,
                    enum_path,
                    variant_name,
                } = parse_enum_path(TypePath { qself, path }).map_err(|path| {
                    state.errors.error(path, "unsupported enum variant path");
                })?;
                Self::enum_variant(
                    state,
                    MatchPatEnumVariant {
                        match_span: state.match_span,
                        variant_path,
                        enum_path,
                        variant_name,
                        field: None,
                    },
                )
            }
            Pat::Struct(PatStruct {
                attrs: _,
                qself,
                path,
                brace_token,
                fields,
                rest,
            }) => {
                let fields = fields
                    .into_pairs()
                    .filter_map_pair_value(|field_pat| {
                        MatchPatStructField::parse(state, field_pat).ok()
                    })
                    .collect();
                if qself.is_some() {
                    state
                        .errors
                        .error(TypePath { qself, path }, "unsupported struct path");
                    return Err(());
                }
                Self::struct_(
                    state,
                    MatchPatStruct {
                        match_span: state.match_span,
                        path,
                        brace_token,
                        fields,
                        rest: rest.map(
                            |PatRest {
                                 attrs: _,
                                 dot2_token,
                             }| dot2_token,
                        ),
                    },
                )
            }
            Pat::TupleStruct(PatTupleStruct {
                attrs: _,
                qself,
                path,
                paren_token,
                mut elems,
            }) => {
                let EnumPath {
                    variant_path,
                    enum_path,
                    variant_name,
                } = parse_enum_path(TypePath { qself, path }).map_err(|path| {
                    state.errors.error(path, "unsupported enum variant path");
                })?;
                if elems.is_empty() {
                    let mut tokens = TokenStream::new();
                    paren_token.surround(&mut tokens, |_| {});
                    state.errors.error(
                        tokens,
                        "field-less enum variants must not be matched using parenthesis",
                    );
                }
                if elems.len() != 1 {
                    state.errors.error(
                        variant_path,
                        "enum variant pattern must have exactly one field",
                    );
                    return Err(());
                }
                let field = elems.pop().unwrap().into_value();
                let field = if let Pat::Rest(rest) = field {
                    MatchPatSimple::Wild(MatchPatWild {
                        underscore_token: Token![_](rest.dot2_token.span()),
                    })
                } else {
                    MatchPatSimple::parse(state, field)?
                };
                Self::enum_variant(
                    state,
                    MatchPatEnumVariant {
                        match_span: state.match_span,
                        variant_path,
                        enum_path,
                        variant_name,
                        field: Some((paren_token, field)),
                    },
                )
            }
            Pat::Rest(_) => {
                state
                    .errors
                    .error(pat, "not allowed here in #[hdl] patterns");
                Err(())
            }
            Pat::Wild(PatWild {
                attrs: _,
                underscore_token,
            }) => Ok(Self::simple(MatchPatSimple::Wild(MatchPatWild {
                underscore_token,
            }))),
            Pat::Tuple(_) | Pat::Slice(_) | Pat::Const(_) | Pat::Lit(_) | Pat::Range(_) => {
                state
                    .errors
                    .error(pat, "not yet implemented in #[hdl] patterns");
                Err(())
            }
            _ => {
                state.errors.error(pat, "not allowed in #[hdl] patterns");
                Err(())
            }
        }
    }
}

impl ParseMatchPat for MatchPatSimple {
    fn simple(v: MatchPatSimple) -> Self {
        v
    }

    fn or(v: MatchPatOr<Self>) -> Self {
        Self::Or(v)
    }

    fn paren(v: MatchPatParen<Self>) -> Self {
        Self::Paren(v)
    }

    fn struct_(state: &mut HdlMatchParseState<'_>, v: MatchPatStruct) -> Result<Self, ()> {
        state.errors.error(
            v.path,
            "matching structs is not yet implemented inside structs/enums in #[hdl] patterns",
        );
        Err(())
    }

    fn enum_variant(
        state: &mut HdlMatchParseState<'_>,
        v: MatchPatEnumVariant,
    ) -> Result<Self, ()> {
        state.errors.error(
            v.variant_path,
            "matching enum variants is not yet implemented inside structs/enums in #[hdl] patterns",
        );
        Err(())
    }
}

#[derive(Debug, Clone)]
enum MatchPat {
    Simple(MatchPatSimple),
    Or(MatchPatOr<MatchPat>),
    Paren(MatchPatParen<MatchPat>),
    Struct(MatchPatStruct),
    EnumVariant(MatchPatEnumVariant),
}

impl_fold! {
    enum MatchPat<> {
        Simple(MatchPatSimple),
        Or(MatchPatOr<MatchPat>),
        Paren(MatchPatParen<MatchPat>),
        Struct(MatchPatStruct),
        EnumVariant(MatchPatEnumVariant),
    }
}

impl ParseMatchPat for MatchPat {
    fn simple(v: MatchPatSimple) -> Self {
        Self::Simple(v)
    }

    fn or(v: MatchPatOr<Self>) -> Self {
        Self::Or(v)
    }

    fn paren(v: MatchPatParen<Self>) -> Self {
        Self::Paren(v)
    }

    fn struct_(_state: &mut HdlMatchParseState<'_>, v: MatchPatStruct) -> Result<Self, ()> {
        Ok(Self::Struct(v))
    }

    fn enum_variant(
        _state: &mut HdlMatchParseState<'_>,
        v: MatchPatEnumVariant,
    ) -> Result<Self, ()> {
        Ok(Self::EnumVariant(v))
    }
}

impl ToTokens for MatchPat {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        match self {
            Self::Simple(v) => v.to_tokens(tokens),
            Self::Or(v) => v.to_tokens(tokens),
            Self::Paren(v) => v.to_tokens(tokens),
            Self::Struct(v) => v.to_tokens(tokens),
            Self::EnumVariant(v) => v.to_tokens(tokens),
        }
    }
}

with_debug_clone_and_fold! {
    struct MatchArm<> {
        attrs: Vec<Attribute>,
        pat: MatchPat,
        fat_arrow_token: Token![=>],
        body: Box<Expr>,
        comma: Option<Token![,]>,
    }
}

impl MatchArm {
    fn parse(state: &mut HdlMatchParseState<'_>, arm: Arm) -> Result<Self, ()> {
        let Arm {
            attrs,
            pat,
            guard,
            fat_arrow_token,
            body,
            comma,
        } = arm;
        if let Some((if_, _)) = guard {
            state
                .errors
                .error(if_, "#[hdl] match arm if clauses are not implemented");
        }
        Ok(Self {
            attrs,
            pat: MatchPat::parse(state, pat)?,
            fat_arrow_token,
            body,
            comma,
        })
    }
}

impl ToTokens for MatchArm {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Self {
            attrs,
            pat,
            fat_arrow_token,
            body,
            comma,
        } = self;
        tokens.append_all(attrs);
        pat.to_tokens(tokens);
        fat_arrow_token.to_tokens(tokens);
        body.to_tokens(tokens);
        comma.to_tokens(tokens);
    }
}

struct RewriteAsCheckMatch {
    span: Span,
}

impl Fold for RewriteAsCheckMatch {
    fn fold_field_pat(&mut self, mut i: FieldPat) -> FieldPat {
        i.colon_token = Some(Token![:](i.member.span()));
        i
    }
    fn fold_pat(&mut self, pat: Pat) -> Pat {
        match pat {
            Pat::Ident(mut pat_ident) => match parse_enum_ident(pat_ident.ident) {
                Ok(EnumPath {
                    variant_path: _,
                    enum_path,
                    variant_name,
                }) => parse_quote_spanned! {self.span=>
                    __MatchTy::<#enum_path>::#variant_name {}
                },
                Err(ident) => {
                    pat_ident.ident = ident;
                    Pat::Ident(self.fold_pat_ident(pat_ident))
                }
            },
            Pat::Path(PatPath {
                attrs: _,
                qself,
                path,
            }) => match parse_enum_path(TypePath { qself, path }) {
                Ok(EnumPath {
                    variant_path: _,
                    enum_path,
                    variant_name,
                }) => parse_quote_spanned! {self.span=>
                    __MatchTy::<#enum_path>::#variant_name {}
                },
                Err(type_path) => parse_quote_spanned! {self.span=>
                    __MatchTy::<#type_path> {}
                },
            },
            Pat::Struct(PatStruct {
                attrs: _,
                qself,
                path,
                brace_token,
                fields,
                rest,
            }) => {
                let type_path = TypePath { qself, path };
                let path = parse_quote_spanned! {self.span=>
                    __MatchTy::<#type_path>
                };
                let fields = fields.do_fold(self);
                Pat::Struct(PatStruct {
                    attrs: vec![],
                    qself: None,
                    path,
                    brace_token,
                    fields,
                    rest,
                })
            }
            Pat::TupleStruct(PatTupleStruct {
                attrs,
                qself,
                path,
                paren_token,
                elems,
            }) => match parse_enum_path(TypePath { qself, path }) {
                Ok(EnumPath {
                    variant_path: _,
                    enum_path,
                    variant_name,
                }) => {
                    let path = parse_quote_spanned! {self.span=>
                        __MatchTy::<#enum_path>::#variant_name
                    };
                    let elems = Punctuated::from_iter(
                        elems.into_pairs().map_pair_value(|p| fold_pat(self, p)),
                    );
                    Pat::TupleStruct(PatTupleStruct {
                        attrs,
                        qself: None,
                        path,
                        paren_token,
                        elems,
                    })
                }
                Err(TypePath { qself, path }) => {
                    Pat::TupleStruct(self.fold_pat_tuple_struct(PatTupleStruct {
                        attrs,
                        qself,
                        path,
                        paren_token,
                        elems,
                    }))
                }
            },
            _ => fold_pat(self, pat),
        }
    }
    fn fold_pat_ident(&mut self, mut i: PatIdent) -> PatIdent {
        i.by_ref = Some(Token![ref](i.ident.span()));
        i.mutability = None;
        i
    }
    fn fold_arm(&mut self, mut i: Arm) -> Arm {
        i.body = parse_quote_spanned! {self.span=>
            match __infallible {}
        };
        i.comma.get_or_insert_with(|| Token![,](self.span));
        fold_arm(self, i)
    }
    fn fold_expr_match(&mut self, mut i: ExprMatch) -> ExprMatch {
        i.expr = parse_quote_spanned! {self.span=>
            __match_value
        };
        fold_expr_match(self, i)
    }
    fn fold_expr(&mut self, i: Expr) -> Expr {
        // don't recurse into expressions
        i
    }
}

struct HdlMatchParseState<'a> {
    match_span: Span,
    errors: &'a mut Errors,
}

impl Visitor<'_> {
    pub(crate) fn process_hdl_match(
        &mut self,
        _hdl_attr: HdlAttr<Nothing>,
        expr_match: ExprMatch,
    ) -> Expr {
        let span = expr_match.match_token.span();
        let check_match = RewriteAsCheckMatch { span }.fold_expr_match(expr_match.clone());
        let ExprMatch {
            attrs: _,
            match_token,
            expr,
            brace_token: _,
            arms,
        } = expr_match;
        self.require_normal_module(match_token);
        let mut state = HdlMatchParseState {
            match_span: span,
            errors: &mut self.errors,
        };
        let arms = Vec::from_iter(
            arms.into_iter()
                .filter_map(|arm| MatchArm::parse(&mut state, arm).ok()),
        );
        let expr = quote_spanned! {span=>
            {
                type __MatchTy<T> = <T as ::fayalite::ty::Type>::MatchVariant;
                let __match_expr = ::fayalite::expr::ToExpr::to_expr(&(#expr));
                ::fayalite::expr::check_match_expr(__match_expr, |__match_value, __infallible| {
                    #[allow(unused_variables)]
                    #check_match
                });
                for __match_variant in ::fayalite::module::match_(__match_expr) {
                    let (__match_variant, __scope) =
                        ::fayalite::ty::MatchVariantAndInactiveScope::match_activate_scope(
                            __match_variant,
                        );
                    #match_token __match_variant {
                        #(#arms)*
                    }
                }
            }
        };
        syn::parse2(expr).unwrap()
    }
}