Skip to main content

strum_macros/macros/
enum_try_as.rs

1use crate::helpers::{case_style::snakify, non_enum_error, HasStrumVariantProperties};
2use proc_macro2::TokenStream;
3use quote::{format_ident, quote, ToTokens};
4use syn::{Data, DeriveInput};
5
6pub fn enum_try_as_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
7    let variants = match &ast.data {
8        Data::Enum(v) => &v.variants,
9        _ => return Err(non_enum_error()),
10    };
11
12    let enum_name = &ast.ident;
13    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
14
15    let variants: Vec<_> = variants
16        .iter()
17        .filter_map(|variant| {
18            if variant.get_variant_properties().ok()?.disabled.is_some() {
19                return None;
20            }
21
22            match &variant.fields {
23                syn::Fields::Unnamed(values) => {
24                    let variant_name = &variant.ident;
25                    let types: Vec<_> = values.unnamed.iter().map(|field| {
26                        field.ty.to_token_stream()
27                    }).collect();
28                    let field_names: Vec<_> = values.unnamed.iter().enumerate().map(|(i, _)| {
29                        let name = "x".repeat(i + 1);
30                        let name = format_ident!("{}", name);
31                        quote! {#name}
32                    }).collect();
33
34                    let move_fn_name = format_ident!("try_as_{}", snakify(&variant_name.to_string()));
35                    let ref_fn_name = format_ident!("try_as_{}_ref", snakify(&variant_name.to_string()));
36                    let mut_fn_name = format_ident!("try_as_{}_mut", snakify(&variant_name.to_string()));
37
38                    Some(quote! {
39                        #[automatically_derived]
40                        #[must_use]
41                        #[inline]
42                        pub fn #move_fn_name(self) -> ::core::option::Option<(#(#types),*)> {
43                            match self {
44                                #enum_name::#variant_name (#(#field_names),*) => Some((#(#field_names),*)),
45                                _ => None
46                            }
47                        }
48
49                        #[automatically_derived]
50                        #[must_use]
51                        #[inline]
52                        pub const fn #ref_fn_name(&self) -> ::core::option::Option<(#(&#types),*)> {
53                            match self {
54                                #enum_name::#variant_name (#(#field_names),*) => Some((#(#field_names),*)),
55                                _ => None
56                            }
57                        }
58
59                        #[automatically_derived]
60                        #[must_use]
61                        #[inline]
62                        pub fn #mut_fn_name(&mut self) -> ::core::option::Option<(#(&mut #types),*)> {
63                            match self {
64                                #enum_name::#variant_name (#(#field_names),*) => Some((#(#field_names),*)),
65                                _ => None
66                            }
67                        }
68                    })
69                },
70                _ => {
71                    None
72                }
73            }
74
75        })
76        .collect();
77
78    Ok(quote! {
79        #[automatically_derived]
80        impl #impl_generics #enum_name #ty_generics #where_clause {
81            #(#variants)*
82        }
83    })
84}