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,
    module::transform_body::{
        expand_aggregate_literals::{AggregateLiteralOptions, StructOrEnumPath},
        with_debug_clone_and_fold, Visitor,
    },
    Errors, HdlAttr,
};
use proc_macro2::{Span, TokenStream};
use quote::{ToTokens, TokenStreamExt};
use syn::{
    fold::{fold_arm, fold_expr_match, fold_pat, Fold},
    parse::Nothing,
    parse_quote_spanned,
    punctuated::{Pair, Punctuated},
    spanned::Spanned,
    token::{Brace, Paren},
    Arm, Attribute, Expr, ExprMatch, FieldPat, Ident, Index, Member, Pat, PatIdent, PatOr,
    PatParen, PatPath, PatRest, PatStruct, PatTupleStruct, PatWild, Path, 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<> {
        member: Member,
        colon_token: Option<Token![:]>,
        pat: MatchPatSimple,
    }
}

impl ToTokens for MatchPatStructField {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Self {
            member,
            colon_token,
            pat,
        } = self;
        member.to_tokens(tokens);
        colon_token.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;
        Ok(Self {
            member,
            colon_token,
            pat: MatchPatSimple::parse(state, *pat)?,
        })
    }
}

with_debug_clone_and_fold! {
    struct MatchPatStruct<> {
        resolved_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 {
            resolved_path,
            brace_token,
            fields,
            rest,
        } = self;
        resolved_path.to_tokens(tokens);
        brace_token.surround(tokens, |tokens| {
            fields.to_tokens(tokens);
            rest.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),
        }
    }
}

fn is_pat_ident_a_struct_or_enum_name(ident: &Ident) -> bool {
    ident
        .to_string()
        .starts_with(|ch: char| ch.is_ascii_uppercase())
}

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,
        struct_error_spanned: &dyn ToTokens,
    ) -> 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");
                }
                if is_pat_ident_a_struct_or_enum_name(&ident) {
                    let ident_span = ident.span();
                    let resolved_path = state.resolve_enum_struct_path(TypePath {
                        qself: None,
                        path: ident.clone().into(),
                    })?;
                    Self::struct_(
                        state,
                        MatchPatStruct {
                            resolved_path,
                            brace_token: Brace(ident_span),
                            fields: Punctuated::new(),
                            rest: None,
                        },
                        &ident,
                    )
                } else {
                    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| {
                        let (pat, punct) = pair.into_tuple();
                        let pat = Self::parse(state, pat).ok()?;
                        Some(Pair::new(pat, punct))
                    })
                    .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 path = TypePath { qself, path };
                let path_span = path.span();
                let resolved_path = state.resolve_enum_struct_path(path.clone())?;
                Self::struct_(
                    state,
                    MatchPatStruct {
                        resolved_path,
                        brace_token: Brace(path_span),
                        fields: Punctuated::new(),
                        rest: None,
                    },
                    &path,
                )
            }
            Pat::Struct(PatStruct {
                attrs: _,
                qself,
                path,
                brace_token,
                fields,
                rest,
            }) => {
                let fields = fields
                    .into_pairs()
                    .filter_map(|pair| {
                        let (field_pat, punct) = pair.into_tuple();
                        let field_pat = MatchPatStructField::parse(state, field_pat).ok()?;
                        Some(Pair::new(field_pat, punct))
                    })
                    .collect();
                let path = TypePath { qself, path };
                let resolved_path = state.resolve_enum_struct_path(path.clone())?;
                Self::struct_(
                    state,
                    MatchPatStruct {
                        resolved_path,
                        brace_token,
                        fields,
                        rest: rest.map(
                            |PatRest {
                                 attrs: _,
                                 dot2_token,
                             }| dot2_token,
                        ),
                    },
                    &path,
                )
            }
            Pat::TupleStruct(PatTupleStruct {
                attrs: _,
                qself,
                path,
                paren_token,
                mut elems,
            }) => {
                let rest = if let Some(&Pat::Rest(PatRest {
                    attrs: _,
                    dot2_token,
                })) = elems.last()
                {
                    elems.pop();
                    Some(dot2_token)
                } else {
                    None
                };
                let fields = elems
                    .into_pairs()
                    .enumerate()
                    .filter_map(|(index, pair)| {
                        let (pat, punct) = pair.into_tuple();
                        let pat = MatchPatSimple::parse(state, pat).ok()?;
                        let mut index = Index::from(index);
                        index.span = state.span;
                        let field = MatchPatStructField {
                            member: index.into(),
                            colon_token: Some(Token![:](state.span)),
                            pat,
                        };
                        Some(Pair::new(field, punct))
                    })
                    .collect();
                let path = TypePath { qself, path };
                let resolved_path = state.resolve_enum_struct_path(path.clone())?;
                Self::struct_(
                    state,
                    MatchPatStruct {
                        resolved_path,
                        brace_token: Brace {
                            span: paren_token.span,
                        },
                        fields,
                        rest,
                    },
                    &path,
                )
            }
            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,
        struct_error_spanned: &dyn ToTokens,
    ) -> Result<Self, ()> {
        state.errors.error(
            struct_error_spanned,
            "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),
}

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

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,
        _struct_error_spanned: &dyn ToTokens,
    ) -> Result<Self, ()> {
        Ok(Self::Struct(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),
        }
    }
}

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, i: Pat) -> Pat {
        match i {
            Pat::Ident(PatIdent {
                attrs,
                by_ref,
                mutability,
                ident,
                subpat: None,
            }) if is_pat_ident_a_struct_or_enum_name(&ident) => {
                parse_quote_spanned! {ident.span()=>
                    #(#attrs)*
                    #by_ref
                    #mutability
                    #ident {}
                }
            }
            _ => fold_pat(self, i),
        }
    }
    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> {
    errors: &'a mut Errors,
    span: Span,
}

impl HdlMatchParseState<'_> {
    fn resolve_enum_struct_path(&mut self, path: TypePath) -> Result<Path, ()> {
        let StructOrEnumPath { ty, variant } =
            StructOrEnumPath::new(self.errors, path, &AggregateLiteralOptions::default())?;
        Ok(if let Some((_variant_path, variant_name)) = variant {
            parse_quote_spanned! {self.span=>
                __MatchTy::<#ty>::#variant_name
            }
        } else {
            parse_quote_spanned! {self.span=>
                __MatchTy::<#ty>
            }
        })
    }
}

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 {
            errors: &mut self.errors,
            span,
        };
        let arms = Vec::from_iter(
            arms.into_iter()
                .filter_map(|arm| MatchArm::parse(&mut state, arm).ok()),
        );
        parse_quote_spanned! {span=>
            {
                type __MatchTy<V> = <<V as ::fayalite::expr::ToExpr>::Type 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 m.match_(__match_expr) {
                    let (__match_variant, __scope) = ::fayalite::ty::MatchVariantAndInactiveScope::match_activate_scope(__match_variant);
                    #match_token __match_variant {
                        #(#arms)*
                    }
                }
            }
        }
    }
}