const_assoc_derive/
lib.rs

1use anyhow::{anyhow, bail, Result};
2use proc_macro::TokenStream;
3use proc_macro2::{Ident, TokenStream as TokenStream2};
4use quote::{quote, ToTokens};
5use syn::{parse_macro_input, AttrStyle, Attribute, Data, DataEnum, DeriveInput};
6
7#[proc_macro_derive(PrimitiveEnum)]
8pub fn derive_primitive_enum(item: TokenStream) -> TokenStream {
9    let input = parse_macro_input!(item as DeriveInput);
10    derive(&input).unwrap().into()
11}
12
13fn derive(input: &DeriveInput) -> Result<TokenStream2> {
14    let data = match &input.data {
15        Data::Enum(data) => data,
16        _ => bail!("`PrimitiveEnum` only applies to enums"),
17    };
18
19    let name = &input.ident;
20    let repr = parse_repr_attribute(&input.attrs)?;
21    let max_variants = max_discriminant(data)? + 1;
22
23    Ok(quote! {
24        unsafe impl ::const_assoc::PrimitiveEnum for #name {
25            type Layout = ::const_assoc::PrimitiveEnumLayout<#repr, #max_variants>;
26        }
27    })
28}
29
30fn parse_repr_attribute(attrs: &[Attribute]) -> Result<EnumRepr> {
31    let attr = attrs
32        .iter()
33        .find(|attr| {
34            let matching_ident = attr
35                .path()
36                .get_ident()
37                .map(|ident| *ident == "repr")
38                .unwrap_or(false);
39
40            matching_ident && matches!(attr.style, AttrStyle::Outer)
41        })
42        .ok_or(anyhow!("enum must have a primitive representation"))?;
43
44    let repr_ident = attr.meta.require_list()?.parse_args::<Ident>()?;
45    let repr_str = repr_ident.to_string();
46
47    let repr = match repr_str.as_str() {
48        "u8" => EnumRepr::U8,
49        "u16" => EnumRepr::U16,
50        "u32" => EnumRepr::U32,
51        "u64" => EnumRepr::U64,
52        "usize" => EnumRepr::USize,
53        _ => bail!("`{repr_str}` is not supported as a primitive enum representation"),
54    };
55
56    Ok(repr)
57}
58
59fn max_discriminant(enum_: &DataEnum) -> Result<usize> {
60    for variant in &enum_.variants {
61        if variant.discriminant.is_some() {
62            bail!("enums that have variants with explicit discriminants are not supported");
63        }
64    }
65
66    Ok(enum_.variants.len() - 1)
67}
68
69enum EnumRepr {
70    U8,
71    U16,
72    U32,
73    U64,
74    USize,
75}
76
77impl ToTokens for EnumRepr {
78    fn to_tokens(&self, tokens: &mut TokenStream2) {
79        let t = match self {
80            EnumRepr::U8 => quote! { u8 },
81            EnumRepr::U16 => quote! { u16 },
82            EnumRepr::U32 => quote! { u32 },
83            EnumRepr::U64 => quote! { u64 },
84            EnumRepr::USize => quote! { usize },
85        };
86        
87        t.to_tokens(tokens);
88    }
89}
const_assoc_derive/lib.rs

const_assoc_derive/
lib.rs
