Skip to main content

assert_order_derive/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, Data, DeriveInput};
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 lifetimes = ast
32        .generics
33        .lifetimes()
34        .map(|lifetime| lifetime.lifetime.clone())
35        .collect::<Vec<_>>();
36
37    let types = ast
38        .generics
39        .type_params()
40        .map(|param| param.ident.clone())
41        .collect::<Vec<_>>();
42
43    let variants = enum_def
44        .variants
45        .iter()
46        .map(|variant| variant.ident.clone())
47        .collect::<Vec<_>>();
48    let variant_count = variants.len();
49
50    let name = &ast.ident;
51    let order_impl = quote! {
52        impl<#(#lifetimes,)* #(#types),*> VariantOrder for #name<#(#lifetimes,)* #(#types),*> {
53            fn order() -> &'static [&'static str] {
54                static VARIANTS: [&'static str; #variant_count] = [
55                    #(stringify!(#variants)),*
56                ];
57
58                return &VARIANTS;
59            }
60        }
61    };
62
63    order_impl.into()
64}