1// SPDX-License-Identifier: Apache-2.0 OR MIT
2
3use core::ops;
4
5use proc_macro2::TokenStream;
6use quote::ToTokens;
7use syn::{
8 parse::{Parse, ParseStream},
9 Fields, Ident, ItemEnum, Result, Type,
10};
11
12/// A structure to make trait implementation to enums more efficient.
13pub struct EnumData {
14 repr: ItemEnum,
15 field_types: Vec<Type>,
16}
17
18impl EnumData {
19 /// Returns an iterator over field types.
20 ///
21 /// ```text
22 /// enum Enum<TypeA, TypeB> {
23 /// VariantA(TypeA),
24 /// ^^^^^
25 /// VariantB(TypeB),
26 /// ^^^^^
27 /// }
28 /// ```
29 pub fn field_types(&self) -> impl ExactSizeIterator<Item = &Type> + Clone {
30 self.field_types.iter()
31 }
32
33 /// Returns an iterator over variant names.
34 ///
35 /// ```text
36 /// enum Enum<TypeA, TypeB> {
37 /// VariantA(TypeA),
38 /// ^^^^^^^^
39 /// VariantB(TypeB),
40 /// ^^^^^^^^
41 /// }
42 /// ```
43 pub fn variant_idents(&self) -> impl ExactSizeIterator<Item = &Ident> + Clone {
44 self.variants.iter().map(|v| &v.ident)
45 }
46}
47
48impl ops::Deref for EnumData {
49 type Target = ItemEnum;
50
51 fn deref(&self) -> &Self::Target {
52 &self.repr
53 }
54}
55
56impl From<EnumData> for ItemEnum {
57 fn from(other: EnumData) -> Self {
58 other.repr
59 }
60}
61
62impl Parse for EnumData {
63 fn parse(input: ParseStream<'_>) -> Result<Self> {
64 let item: ItemEnum = input.parse()?;
65
66 if item.variants.is_empty() {
67 bail!(item, "may not be used on enums without variants");
68 }
69
70 let field_types = item.variants.iter().try_fold(
71 Vec::with_capacity(item.variants.len()),
72 |mut field_types, v| {
73 if let Some((_, e)) = &v.discriminant {
74 bail!(e, "may not be used on enums with discriminants");
75 }
76
77 if v.fields.is_empty() {
78 bail!(v, "may not be used on enums with variants with zero fields");
79 } else if v.fields.len() != 1 {
80 bail!(v, "may not be used on enums with variants with multiple fields");
81 }
82
83 match &v.fields {
84 Fields::Unnamed(f) => {
85 field_types.push(f.unnamed.iter().next().unwrap().ty.clone());
86 Ok(field_types)
87 }
88 Fields::Named(_) => {
89 bail!(v, "may not be used on enums with variants with named fields");
90 }
91 Fields::Unit => unreachable!(),
92 }
93 },
94 )?;
95
96 Ok(Self { repr: item, field_types })
97 }
98}
99
100impl ToTokens for EnumData {
101 fn to_tokens(&self, tokens: &mut TokenStream) {
102 self.repr.to_tokens(tokens);
103 }
104}
105