cenum_utils_derive/
lib.rs

1/*!
2[![github]](https://github.com/Josh194/ConstEnumUtils) [![crates-io]](https://crates.io/crates/cenum-utils-derive) [![docs-rs]](https://docs.rs/cenum-utils-derive)
3
4[github]: https://img.shields.io/badge/github-8da0cb?style=for-the-badge&labelColor=555555&logo=github
5[crates-io]: https://img.shields.io/badge/crates.io-fc8d62?style=for-the-badge&labelColor=555555&logo=rust
6[docs-rs]: https://img.shields.io/badge/docs.rs-66c2a5?style=for-the-badge&labelColor=555555&logo=docs.rs
7
8A helper crate providing derive macros for `cenum_utils`.
9
10See the parent docs for more information on use.
11
12# Example
13
14```rust
15use cenum_utils_derive::ConstEnum;
16
17#[derive(ConstEnum)]
18#[repr(u8)]
19enum Enum {
20	X,
21	Y,
22	Z
23}
24```
25*/
26
27use std::{collections::HashSet};
28
29use proc_macro2::{Span, TokenStream};
30
31use quote::quote;
32use rustc_hash::{FxBuildHasher};
33use syn::{parse_macro_input, DeriveInput, Ident, Type};
34
35/// Derives `EnumCount` and `EnumNames`, as well as `EnumDiscriminants` if the enum has a valid primitive `repr` type.
36/// 
37/// Valid `repr` types are `u8`, `u16`, `u32`, `u64`, `u128`, `usize`, `i8`, `i16`, `i32`, `i64`, `i128`, and `isize`.
38#[proc_macro_derive(ConstEnum)]
39pub fn derive_const_enum(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
40	impl_const_enum(parse_macro_input!(input as DeriveInput)).unwrap_or_else(|error| error.into_compile_error()).into()
41}
42
43thread_local! {
44	/// The set of all valid primitive enum representations.
45	static PRIMITIVE_TYPES: HashSet<Ident, FxBuildHasher> = {
46		let type_tokens = [
47			quote! { u8 },
48			quote! { u16 },
49			quote! { u32 },
50			quote! { u64 },
51			quote! { u128 },
52			quote! { usize },
53			quote! { i8 },
54			quote! { i16 },
55			quote! { i32 },
56			quote! { i64 },
57			quote! { i128 },
58			quote! { isize },
59		];
60
61		HashSet::from_iter(type_tokens.into_iter().map(|tokens| syn::parse2(tokens).unwrap()))
62	};
63}
64
65fn impl_const_enum(item: DeriveInput) -> Result<TokenStream, syn::Error> {
66	let syn::Data::Enum(enum_data) = item.data else { return Err(syn::Error::new_spanned(item, "ConstEnum: expected an enum")); };
67
68	// * Get the enum's `repr` type if it exists and is of the expected format.
69	let representation: Option<Type> = item.attrs.iter().find_map(|attribute| {
70		attribute.meta
71			.path()
72			.get_ident()
73			.and_then(|ident| {
74				(ident == "repr")
75					.then(|| attribute.parse_args().ok())
76					.flatten()
77			})
78	});
79
80	let enum_name: Ident = item.ident;
81	let variant_count: usize = enum_data.variants.len();
82
83	let variant_idents = enum_data.variants.iter().map(|variant| &variant.ident);
84	let variant_names = enum_data.variants.iter().map(|variant| variant.ident.to_string());
85
86	// * Get the representation type if it is a valid primitive.
87	let enum_discriminant: Option<Type> = representation.and_then(|ty| {
88		let Type::Path(type_path) = &ty else { return None; };
89
90		if !type_path.qself.is_none() { return None; };
91		let ident: &Ident = type_path.path.get_ident()?;
92
93		PRIMITIVE_TYPES.with(|primitives| primitives.contains(ident)).then_some(ty)
94	});
95
96	let crate_name: Ident = Ident::new("cenum_utils", Span::call_site());
97
98	let count_impl: TokenStream = quote! {
99		impl ::#crate_name::EnumCount for #enum_name {
100			const COUNT: usize = #variant_count;
101		}
102	};
103
104	let names_impl: TokenStream = quote! {
105		impl ::#crate_name::EnumNames for #enum_name {
106			const NAMES: &[&str] = &[#(#variant_names),*];
107		}
108	};
109
110	let discriminants_impl: Option<TokenStream> = enum_discriminant.map(|discriminant| {
111		quote! {
112			impl ::#crate_name::EnumDiscriminants for #enum_name {
113				type Discriminant = #discriminant;
114
115				const DISCRIMINANTS: &[Self::Discriminant] = &[#(Self::#variant_idents as #discriminant),*];
116			}
117		}
118	});
119
120	Ok(quote! { #count_impl #names_impl #discriminants_impl })
121}