Skip to main content

assert_order_derive/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{Data, DeriveInput, parse_macro_input};
4
5/// Derive macro for implementing `VariantOrder` based on
6/// the definition of an enum.
7#[proc_macro_derive(VariantOrder)]
8pub fn variant_order_derive(input: TokenStream) -> TokenStream {
9    let ast = parse_macro_input!(input as DeriveInput);
10
11    let enum_def = match ast.data {
12        Data::Enum(ref enum_def) => enum_def,
13        Data::Union(union_def) => {
14            return syn::Error::new_spanned(
15                &union_def.union_token,
16                "VariantOrder can only be applied to enums.",
17            )
18            .to_compile_error()
19            .into();
20        }
21        Data::Struct(struct_def) => {
22            return syn::Error::new_spanned(
23                &struct_def.struct_token,
24                "VariantOrder can only be applied to enums.",
25            )
26            .to_compile_error()
27            .into();
28        }
29    };
30
31    let variants = enum_def
32        .variants
33        .iter()
34        .map(|variant| variant.ident.clone())
35        .collect::<Vec<_>>();
36    let variant_count = variants.len();
37
38    let name = &ast.ident;
39    let order_impl = quote! {
40        impl VariantOrder for #name {
41            fn order() -> &'static [&'static str] {
42                static VARIANTS: [&'static str; #variant_count] = [
43                    #(stringify!(#variants)),*
44                ];
45
46                return &VARIANTS;
47            }
48        }
49    };
50
51    order_impl.into()
52}