unit_enum/lib.rs
1#![doc = include_str!("lib.md")]
2
3use proc_macro::TokenStream;
4use quote::quote;
5use syn::punctuated::Punctuated;
6use syn::{parse_macro_input, Data, DeriveInput, Error, Expr, Fields, Meta, Token, Type, Variant};
7
8// Compile and run the README examples as doctests so they stay in sync with the crate
9#[cfg(doctest)]
10#[doc = include_str!("../README.md")]
11struct ReadmeDoctests;
12
13/// Derives the `UnitEnum` trait for an enum.
14///
15/// This macro can be used on enums with unit variants (no fields) and optionally one "other" variant
16/// that can hold arbitrary discriminant values.
17///
18/// # Attributes
19/// - `#[repr(type)]`: Optional for regular enums, defaults to i32. Required when using an "other" variant.
20/// - `#[unit_enum(other)]`: Marks a variant as the catch-all for undefined discriminant values.
21/// The type of this variant must match the repr type.
22///
23/// # Requirements
24/// - The enum must contain only unit variants, except for one optional "other" variant
25/// - The "other" variant, if present, must:
26/// - Be marked with `#[unit_enum(other)]`
27/// - Have exactly one unnamed field matching the repr type
28/// - Be the only variant with the "other" attribute
29/// - Have a matching `#[repr(type)]` attribute
30///
31/// # Examples
32///
33/// Basic usage with unit variants (repr is optional):
34/// ```rust
35/// # use unit_enum::UnitEnum;
36/// #[derive(UnitEnum)]
37/// enum Example {
38/// A,
39/// B = 10,
40/// C,
41/// }
42/// ```
43///
44/// Usage with explicit repr:
45/// ```rust
46/// # use unit_enum::UnitEnum;
47/// #[derive(UnitEnum)]
48/// #[repr(u16)]
49/// enum Color {
50/// Red = 10,
51/// Green,
52/// Blue = 45654,
53/// }
54/// ```
55///
56/// Usage with an "other" variant (repr required):
57/// ```rust
58/// # use unit_enum::UnitEnum;
59/// #[derive(UnitEnum)]
60/// #[repr(u16)]
61/// enum Status {
62/// Active = 1,
63/// Inactive = 2,
64/// #[unit_enum(other)]
65/// Unknown(u16), // type must match repr
66/// }
67/// ```
68#[proc_macro_derive(UnitEnum, attributes(unit_enum))]
69pub fn unit_enum_derive(input: TokenStream) -> TokenStream {
70 let ast = parse_macro_input!(input as DeriveInput);
71
72 match validate_and_process(&ast) {
73 Ok(validation) => impl_unit_enum(
74 &ast,
75 &validation.discriminant_type,
76 &validation.unit_variants,
77 validation.other_variant,
78 ),
79 Err(e) => e.to_compile_error().into(),
80 }
81}
82
83struct ValidationResult<'a> {
84 discriminant_type: Type,
85 unit_variants: Vec<&'a Variant>,
86 other_variant: Option<(&'a Variant, Type)>,
87}
88
89fn validate_and_process(ast: &DeriveInput) -> Result<ValidationResult<'_>, Error> {
90 let data_enum = match &ast.data {
91 Data::Enum(data_enum) => data_enum,
92 _ => return Err(Error::new_spanned(ast, "UnitEnum can only be derived for enums")),
93 };
94
95 let mut validation = ValidationResult {
96 // Get discriminant type from #[repr] attribute
97 discriminant_type: get_discriminant_type(ast)?,
98 unit_variants: Vec::new(),
99 other_variant: None,
100 };
101
102 // Validate each variant
103 for variant in &data_enum.variants {
104 match &variant.fields {
105 Fields::Unit => {
106 if has_unit_enum_attr(variant) {
107 return Err(Error::new_spanned(variant,
108 "Unit variants cannot have #[unit_enum] attributes"));
109 }
110 validation.unit_variants.push(variant);
111 }
112 Fields::Unnamed(fields) if fields.unnamed.len() == 1 => {
113 if has_unit_enum_other_attr(variant) {
114 if validation.other_variant.is_some() {
115 return Err(Error::new_spanned(variant,
116 "Multiple #[unit_enum(other)] variants found. Only one is allowed"));
117 }
118 validation.other_variant = Some((variant, fields.unnamed[0].ty.clone()));
119 } else {
120 return Err(Error::new_spanned(variant,
121 "Non-unit variant must be marked with #[unit_enum(other)] to be used as the catch-all variant"));
122 }
123 }
124 _ => return Err(Error::new_spanned(variant,
125 "Invalid variant. UnitEnum only supports unit variants and a single tuple variant marked with #[unit_enum(other)]")),
126 }
127 }
128
129 Ok(validation)
130}
131
132fn get_discriminant_type(ast: &DeriveInput) -> Result<Type, Error> {
133 const INTEGER_TYPES: [&str; 12] = [
134 "i8", "i16", "i32", "i64", "i128", "isize", "u8", "u16", "u32", "u64", "u128", "usize",
135 ];
136
137 for attr in ast.attrs.iter().filter(|attr| attr.path().is_ident("repr")) {
138 let hints = attr
139 .parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
140 .map_err(|_| Error::new_spanned(attr, "Invalid repr attribute"))?;
141
142 // Hints such as `C` don't set the discriminant type, so look for an integer type among them
143 for hint in hints {
144 if let Meta::Path(path) = hint {
145 if INTEGER_TYPES.iter().any(|ty| path.is_ident(ty)) {
146 return Ok(syn::parse_quote!(#path));
147 }
148 }
149 }
150 }
151
152 Ok(syn::parse_quote!(i32))
153}
154
155fn has_unit_enum_attr(variant: &Variant) -> bool {
156 variant.attrs.iter().any(|attr| attr.path().is_ident("unit_enum"))
157}
158
159fn has_unit_enum_other_attr(variant: &Variant) -> bool {
160 variant.attrs.iter().any(|attr| {
161 attr.path().is_ident("unit_enum")
162 && attr
163 .parse_nested_meta(|meta| {
164 if meta.path.is_ident("other") {
165 Ok(())
166 } else {
167 Err(meta.error("Invalid unit_enum attribute"))
168 }
169 })
170 .is_ok()
171 })
172}
173
174fn compute_discriminants(variants: &[&Variant]) -> Vec<Expr> {
175 let mut discriminants = Vec::with_capacity(variants.len());
176 // Last explicit discriminant expression, and how many implicit variants followed it
177 let mut base: Option<&Expr> = None;
178 let mut offset: usize = 0;
179
180 for variant in variants {
181 let discriminant = if let Some((_, expr)) = &variant.discriminant {
182 base = Some(expr);
183 offset = 0;
184 expr.clone()
185 } else if let Some(expr) = base {
186 offset += 1;
187 // Parenthesize the base so operators like `<<` or `|` keep their meaning: `(1 << 3) + 1`
188 let offset = proc_macro2::Literal::usize_unsuffixed(offset);
189 syn::parse_quote! { (#expr) + #offset }
190 } else {
191 // No explicit discriminant yet: variants count up from 0
192 let value = proc_macro2::Literal::usize_unsuffixed(offset);
193 offset += 1;
194 syn::parse_quote! { #value }
195 };
196
197 discriminants.push(discriminant);
198 }
199
200 discriminants
201}
202
203fn impl_unit_enum(
204 ast: &DeriveInput, discriminant_type: &Type, unit_variants: &[&Variant], other_variant: Option<(&Variant, Type)>,
205) -> TokenStream {
206 let name = &ast.ident;
207 let num_variants = unit_variants.len();
208 let discriminants = compute_discriminants(unit_variants);
209
210 let name_impl = generate_name_impl(name, unit_variants, &other_variant);
211 let ordinal_impl = generate_ordinal_impl(name, unit_variants, &other_variant, num_variants);
212 let from_ordinal_impl = generate_from_ordinal_impl(name, unit_variants);
213 let discriminant_impl =
214 generate_discriminant_impl(name, unit_variants, &other_variant, discriminant_type, &discriminants);
215 let from_discriminant_impl =
216 generate_from_discriminant_impl(name, unit_variants, &other_variant, discriminant_type, &discriminants);
217 let values_impl = generate_values_impl(name, unit_variants, num_variants);
218
219 quote! {
220 impl #name {
221 #name_impl
222
223 #ordinal_impl
224
225 #from_ordinal_impl
226
227 #discriminant_impl
228
229 #from_discriminant_impl
230
231 /// Returns the total number of unit variants in the enum (excluding the "other" variant if present).
232 ///
233 /// # Examples
234 ///
235 /// ```ignore
236 /// # use unit_enum::UnitEnum;
237 /// #[derive(UnitEnum)]
238 /// enum Example {
239 /// A,
240 /// B,
241 /// #[unit_enum(other)]
242 /// Other(i32),
243 /// }
244 ///
245 /// assert_eq!(Example::len(), 2);
246 /// ```
247 pub const fn len() -> usize {
248 #num_variants
249 }
250
251 #values_impl
252 }
253 }
254 .into()
255}
256
257fn generate_name_impl(
258 name: &syn::Ident, unit_variants: &[&Variant], other_variant: &Option<(&Variant, Type)>,
259) -> proc_macro2::TokenStream {
260 let unit_match_arms = unit_variants.iter().map(|variant| {
261 let variant_name = &variant.ident;
262 quote! { #name::#variant_name => ::core::stringify!(#variant_name) }
263 });
264
265 let other_arm = other_variant.as_ref().map(|(variant, _)| {
266 let variant_name = &variant.ident;
267 quote! { #name::#variant_name(_) => ::core::stringify!(#variant_name) }
268 });
269
270 quote! {
271 /// Returns the name of the enum variant as a string.
272 ///
273 /// # Examples
274 ///
275 /// ```ignore
276 /// # use unit_enum::UnitEnum;
277 /// #[derive(UnitEnum)]
278 /// enum Example {
279 /// A,
280 /// B = 10,
281 /// C,
282 /// }
283 ///
284 /// assert_eq!(Example::A.name(), "A");
285 /// assert_eq!(Example::B.name(), "B");
286 /// assert_eq!(Example::C.name(), "C");
287 /// ```
288 pub const fn name(&self) -> &'static str {
289 match self {
290 #(#unit_match_arms,)*
291 #other_arm
292 }
293 }
294 }
295}
296
297fn generate_ordinal_impl(
298 name: &syn::Ident, unit_variants: &[&Variant], other_variant: &Option<(&Variant, Type)>, num_variants: usize,
299) -> proc_macro2::TokenStream {
300 let unit_match_arms = unit_variants.iter().enumerate().map(|(index, variant)| {
301 let variant_name = &variant.ident;
302 quote! { #name::#variant_name => #index }
303 });
304
305 let other_arm = other_variant.as_ref().map(|(variant, _)| {
306 let variant_name = &variant.ident;
307 quote! { #name::#variant_name(_) => #num_variants }
308 });
309
310 quote! {
311 /// Returns the zero-based ordinal of the enum variant.
312 ///
313 /// For enums with an "other" variant, it returns the position after all unit variants.
314 ///
315 /// # Examples
316 ///
317 /// ```ignore
318 /// # use unit_enum::UnitEnum;
319 /// #[derive(UnitEnum)]
320 /// enum Example {
321 /// A, // ordinal: 0
322 /// B = 10, // ordinal: 1
323 /// C, // ordinal: 2
324 /// }
325 ///
326 /// assert_eq!(Example::A.ordinal(), 0);
327 /// assert_eq!(Example::B.ordinal(), 1);
328 /// assert_eq!(Example::C.ordinal(), 2);
329 /// ```
330 pub const fn ordinal(&self) -> usize {
331 match self {
332 #(#unit_match_arms,)*
333 #other_arm
334 }
335 }
336 }
337}
338fn generate_from_ordinal_impl(name: &syn::Ident, unit_variants: &[&Variant]) -> proc_macro2::TokenStream {
339 let match_arms = unit_variants.iter().enumerate().map(|(index, variant)| {
340 let variant_name = &variant.ident;
341 quote! { #index => ::core::option::Option::Some(#name::#variant_name) }
342 });
343
344 quote! {
345 /// Converts a zero-based ordinal to an enum variant, if possible.
346 ///
347 /// Returns `Some(variant)` if the ordinal corresponds to a unit variant,
348 /// or `None` if the ordinal is out of range or would correspond to the "other" variant.
349 ///
350 /// # Examples
351 ///
352 /// ```ignore
353 /// # use unit_enum::UnitEnum;
354 /// # #[derive(Debug, PartialEq)]
355 /// #[derive(UnitEnum)]
356 /// enum Example {
357 /// A,
358 /// B,
359 /// #[unit_enum(other)]
360 /// Other(i32),
361 /// }
362 ///
363 /// assert_eq!(Example::from_ordinal(0), Some(Example::A));
364 /// assert_eq!(Example::from_ordinal(2), None); // Other variant
365 /// assert_eq!(Example::from_ordinal(99), None); // Out of range
366 /// ```
367 pub const fn from_ordinal(ord: usize) -> ::core::option::Option<Self> {
368 match ord {
369 #(#match_arms,)*
370 _ => ::core::option::Option::None
371 }
372 }
373 }
374}
375
376fn generate_discriminant_impl(
377 name: &syn::Ident, unit_variants: &[&Variant], other_variant: &Option<(&Variant, Type)>, discriminant_type: &Type,
378 discriminants: &[Expr],
379) -> proc_macro2::TokenStream {
380 let unit_match_arms = unit_variants.iter().zip(discriminants).map(|(variant, discriminant)| {
381 let variant_name = &variant.ident;
382 quote! { #name::#variant_name => #discriminant }
383 });
384
385 let other_arm = other_variant.as_ref().map(|(variant, _)| {
386 let variant_name = &variant.ident;
387 quote! { #name::#variant_name(val) => *val }
388 });
389
390 quote! {
391 /// Returns the discriminant value of the enum variant.
392 ///
393 /// For "other" variants, returns the contained value.
394 ///
395 /// # Examples
396 ///
397 /// ```ignore
398 /// # use unit_enum::UnitEnum;
399 /// #[derive(UnitEnum)]
400 /// enum Example {
401 /// A, // 0
402 /// B = 10, // 10
403 /// C, // 11
404 /// }
405 ///
406 /// assert_eq!(Example::A.discriminant(), 0);
407 /// assert_eq!(Example::B.discriminant(), 10);
408 /// assert_eq!(Example::C.discriminant(), 11);
409 /// ```
410 pub const fn discriminant(&self) -> #discriminant_type {
411 match self {
412 #(#unit_match_arms,)*
413 #other_arm
414 }
415 }
416 }
417}
418
419fn generate_from_discriminant_impl(
420 name: &syn::Ident, unit_variants: &[&Variant], other_variant: &Option<(&Variant, Type)>, discriminant_type: &Type,
421 discriminants: &[Expr],
422) -> proc_macro2::TokenStream {
423 if let Some((other_variant, _)) = other_variant {
424 let match_arms = unit_variants.iter().zip(discriminants).map(|(variant, discriminant)| {
425 let variant_name = &variant.ident;
426 quote! { x if x == #discriminant => #name::#variant_name }
427 });
428
429 let other_name = &other_variant.ident;
430 quote! {
431 /// Converts a discriminant value to an enum variant.
432 ///
433 /// For enums with an "other" variant, this will always return a value,
434 /// using the "other" variant for undefined discriminants.
435 ///
436 /// # Examples
437 ///
438 /// ```ignore
439 /// # use unit_enum::UnitEnum;
440 /// #[derive(UnitEnum, PartialEq, Debug)]
441 /// #[repr(u8)]
442 /// enum Example {
443 /// A, // 0
444 /// B = 10, // 10
445 /// #[unit_enum(other)]
446 /// Other(u8),
447 /// }
448 ///
449 /// assert_eq!(Example::from_discriminant(0), Example::A);
450 /// assert_eq!(Example::from_discriminant(10), Example::B);
451 /// assert_eq!(Example::from_discriminant(42), Example::Other(42));
452 /// ```
453 pub const fn from_discriminant(discr: #discriminant_type) -> Self {
454 match discr {
455 #(#match_arms,)*
456 other => #name::#other_name(other)
457 }
458 }
459 }
460 } else {
461 let match_arms = unit_variants.iter().zip(discriminants).map(|(variant, discriminant)| {
462 let variant_name = &variant.ident;
463 quote! { x if x == #discriminant => ::core::option::Option::Some(#name::#variant_name) }
464 });
465
466 quote! {
467 /// Converts a discriminant value to an enum variant, if possible.
468 ///
469 /// Returns `Some(variant)` if the discriminant corresponds to a defined variant,
470 /// or `None` if the discriminant is undefined.
471 ///
472 /// # Examples
473 ///
474 /// ```ignore
475 /// # use unit_enum::UnitEnum;
476 /// #[derive(UnitEnum, PartialEq, Debug)]
477 /// #[repr(u8)]
478 /// enum Example {
479 /// A, // 0
480 /// B = 10, // 10
481 /// C, // 11
482 /// }
483 ///
484 /// assert_eq!(Example::from_discriminant(0), Some(Example::A));
485 /// assert_eq!(Example::from_discriminant(10), Some(Example::B));
486 /// assert_eq!(Example::from_discriminant(42), None);
487 /// ```
488 pub const fn from_discriminant(discr: #discriminant_type) -> ::core::option::Option<Self> {
489 match discr {
490 #(#match_arms,)*
491 _ => ::core::option::Option::None
492 }
493 }
494 }
495 }
496}
497
498fn generate_values_impl(
499 name: &syn::Ident, unit_variants: &[&Variant], num_variants: usize,
500) -> proc_macro2::TokenStream {
501 let variant_exprs = unit_variants.iter().map(|variant| {
502 let variant_name = &variant.ident;
503 quote! { #name::#variant_name }
504 });
505
506 quote! {
507 /// All unit variants of the enum, in declaration order.
508 ///
509 /// Note: This does not include the "other" variant, if present.
510 ///
511 /// # Examples
512 ///
513 /// ```ignore
514 /// # use unit_enum::UnitEnum;
515 /// #[derive(UnitEnum, PartialEq, Debug)]
516 /// enum Example {
517 /// A,
518 /// B,
519 /// #[unit_enum(other)]
520 /// Other(i32),
521 /// }
522 ///
523 /// const ALL: [Example; Example::len()] = Example::VALUES;
524 /// assert_eq!(ALL, [Example::A, Example::B]);
525 /// ```
526 pub const VALUES: [Self; #num_variants] = [
527 #(#variant_exprs),*
528 ];
529
530 /// Returns an iterator over all unit variants of the enum.
531 ///
532 /// Note: This does not include values from the "other" variant, if present.
533 ///
534 /// # Examples
535 ///
536 /// ```ignore
537 /// # use unit_enum::UnitEnum;
538 /// #[derive(UnitEnum, PartialEq, Debug)]
539 /// enum Example {
540 /// A,
541 /// B,
542 /// #[unit_enum(other)]
543 /// Other(i32),
544 /// }
545 ///
546 /// let values: Vec<_> = Example::values().collect();
547 /// assert_eq!(values, vec![Example::A, Example::B]);
548 /// ```
549 pub fn values() -> impl ::core::iter::Iterator<Item = Self> {
550 ::core::iter::IntoIterator::into_iter(Self::VALUES)
551 }
552 }
553}