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