strum_macros/macros/
from_repr.rs

1use heck::ToShoutySnakeCase;
2use proc_macro2::{Span, TokenStream};
3use quote::{format_ident, quote, ToTokens};
4use syn::{Data, DeriveInput, Fields, PathArguments, Type, TypeParen};
5
6use crate::helpers::{non_enum_error, HasStrumVariantProperties};
7
8pub fn from_repr_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
9    let name = &ast.ident;
10    let gen = &ast.generics;
11    let (impl_generics, ty_generics, where_clause) = gen.split_for_impl();
12    let vis = &ast.vis;
13    let attrs = &ast.attrs;
14
15    let mut discriminant_type: Type = syn::parse("usize".parse().unwrap()).unwrap();
16    for attr in attrs {
17        let path = attr.path();
18
19        let mut ts = if let Ok(ts) = attr
20            .meta
21            .require_list()
22            .map(|metas| metas.to_token_stream().into_iter())
23        {
24            ts
25        } else {
26            continue;
27        };
28        // Discard the path
29        let _ = ts.next();
30        let tokens: TokenStream = ts.collect();
31
32        if path.leading_colon.is_some() {
33            continue;
34        }
35        if path.segments.len() != 1 {
36            continue;
37        }
38        let segment = path.segments.first().unwrap();
39        if segment.ident != "repr" {
40            continue;
41        }
42        if segment.arguments != PathArguments::None {
43            continue;
44        }
45        let typ_paren = match syn::parse2::<Type>(tokens.clone()) {
46            Ok(Type::Paren(TypeParen { elem, .. })) => *elem,
47            _ => continue,
48        };
49        let inner_path = match &typ_paren {
50            Type::Path(t) => t,
51            _ => continue,
52        };
53        if let Some(seg) = inner_path.path.segments.last() {
54            for t in &[
55                "u8", "u16", "u32", "u64", "usize", "i8", "i16", "i32", "i64", "isize",
56            ] {
57                if seg.ident == t {
58                    discriminant_type = typ_paren;
59                    break;
60                }
61            }
62        }
63    }
64
65    if gen.lifetimes().count() > 0 {
66        return Err(syn::Error::new(
67            Span::call_site(),
68            "This macro doesn't support enums with lifetimes. \
69             The resulting enums would be unbounded.",
70        ));
71    }
72
73    let variants = match &ast.data {
74        Data::Enum(v) => &v.variants,
75        _ => return Err(non_enum_error()),
76    };
77
78    let mut arms = Vec::new();
79    let mut constant_defs = Vec::new();
80    let mut has_additional_data = false;
81    let mut prev_const_var_ident = None;
82    for variant in variants {
83        if variant.get_variant_properties()?.disabled.is_some() {
84            continue;
85        }
86
87        let ident = &variant.ident;
88        let params = match &variant.fields {
89            Fields::Unit => quote! {},
90            Fields::Unnamed(fields) => {
91                has_additional_data = true;
92                let defaults = ::core::iter::repeat(quote!(::core::default::Default::default()))
93                    .take(fields.unnamed.len());
94                quote! { (#(#defaults),*) }
95            }
96            Fields::Named(fields) => {
97                has_additional_data = true;
98                let fields = fields
99                    .named
100                    .iter()
101                    .map(|field| field.ident.as_ref().unwrap());
102                quote! { {#(#fields: ::core::default::Default::default()),*} }
103            }
104        };
105
106        let const_var_str = format!("{}_DISCRIMINANT", variant.ident).to_shouty_snake_case();
107        let const_var_ident = format_ident!("{}", const_var_str);
108
109        let const_val_expr = match &variant.discriminant {
110            Some((_, expr)) => quote! { #expr },
111            None => match &prev_const_var_ident {
112                Some(prev) => quote! { #prev + 1 },
113                None => quote! { 0 },
114            },
115        };
116
117        constant_defs.push(quote! {const #const_var_ident: #discriminant_type = #const_val_expr;});
118        arms.push(quote! {v if v == #const_var_ident => ::core::option::Option::Some(#name::#ident #params)});
119
120        prev_const_var_ident = Some(const_var_ident);
121    }
122
123    arms.push(quote! { _ => ::core::option::Option::None });
124
125    let const_if_possible = if has_additional_data {
126        quote! {}
127    } else {
128        #[rustversion::before(1.46)]
129        fn filter_by_rust_version(_: TokenStream) -> TokenStream {
130            quote! {}
131        }
132
133        #[rustversion::since(1.46)]
134        fn filter_by_rust_version(s: TokenStream) -> TokenStream {
135            s
136        }
137        filter_by_rust_version(quote! { const })
138    };
139
140    Ok(quote! {
141        #[allow(clippy::use_self)]
142        impl #impl_generics #name #ty_generics #where_clause {
143            #[doc = "Try to create [Self] from the raw representation"]
144            #vis #const_if_possible fn from_repr(discriminant: #discriminant_type) -> Option<#name #ty_generics> {
145                #(#constant_defs)*
146                match discriminant {
147                    #(#arms),*
148                }
149            }
150        }
151    })
152}