repr_fits/
lib.rs

1//! Compile-time bit-width checks for enum discriminants.
2//!
3//! `repr_fits` provides an attribute macro for enums whose discriminants are packed into a fixed
4//! number of bits in another representation, such as a protocol field or compact integer id.
5
6use proc_macro::TokenStream;
7use quote::quote;
8use syn::parse::{Parse, ParseStream};
9use syn::{Error, Fields, Ident, ItemEnum, LitInt, Result, Token, parse_macro_input};
10
11struct ReprFitsArgs {
12    bits: u32,
13}
14
15impl Parse for ReprFitsArgs {
16    fn parse(input: ParseStream<'_>) -> Result<Self> {
17        let name = input.parse::<Ident>()?;
18        if name != "bits" {
19            return Err(Error::new(name.span(), "expected `bits = N`"));
20        }
21
22        input.parse::<Token![=]>()?;
23        let bits = input.parse::<LitInt>()?.base10_parse::<u32>()?;
24        if bits >= 128 {
25            return Err(Error::new(name.span(), "`bits` must be less than 128"));
26        }
27
28        Ok(Self { bits })
29    }
30}
31
32/// Assert at compile time that each enum variant discriminant fits in `bits` bits.
33///
34/// The macro preserves the enum definition and appends a private `const` assertion block. It is
35/// intended for fieldless enums with primitive integer representations such as `#[repr(u8)]`.
36///
37/// # Example
38///
39/// ```rust
40/// use repr_fits::repr_fits;
41///
42/// #[repr_fits(bits = 5)]
43/// #[repr(u8)]
44/// enum RegionCode {
45///     Local = 0,
46///     Backup = 4,
47/// }
48/// ```
49///
50/// A variant outside the declared bit width fails at compile time:
51///
52/// ```compile_fail
53/// use repr_fits::repr_fits;
54///
55/// #[repr_fits(bits = 2)]
56/// #[repr(u8)]
57/// enum PacketKind {
58///     Data = 0,
59///     Ack = 1,
60///     Control = 4,
61/// }
62/// ```
63#[proc_macro_attribute]
64pub fn repr_fits(args: TokenStream, input: TokenStream) -> TokenStream {
65    let args = parse_macro_input!(args as ReprFitsArgs);
66    let item = parse_macro_input!(input as ItemEnum);
67
68    let enum_ident = &item.ident;
69    let mut variant_idents = Vec::with_capacity(item.variants.len());
70    let mut errors = Vec::new();
71
72    for variant in &item.variants {
73        match variant.fields {
74            Fields::Unit => variant_idents.push(&variant.ident),
75            _ => errors.push(Error::new_spanned(variant, "`repr_fits` only supports unit enum variants")),
76        }
77    }
78
79    if !errors.is_empty() {
80        let compile_errors = errors.into_iter().map(Error::into_compile_error);
81        return quote! {
82            #item
83            #(#compile_errors)*
84        }
85        .into();
86    }
87
88    let bits = args.bits;
89    let assertions = variant_idents.iter().map(|variant_ident| {
90        let message = format!("{enum_ident}::{variant_ident} discriminant does not fit in {bits} bits");
91        quote! {
92            assert!((#enum_ident::#variant_ident as u128) < (1u128 << #bits), #message);
93        }
94    });
95
96    quote! {
97        #item
98
99        const _: () = {
100            #(#assertions)*
101        };
102    }
103    .into()
104}
repr_fits/lib.rs

repr_fits/
lib.rs
