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