derive_utils/
ast.rs

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}
derive_utils/ast.rs

derive_utils/
ast.rs
