1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

/*!
This crate provides `VariantCount` derive macro for enum which adds to it the `VARIANT_COUNT` constant, containing count of enum variants.

## Example

```rust
extern crate variant_count;

use variant_count::VariantCount;

#[derive(VariantCount)]
enum Test {
    First(i32),
    Second(Option<String>),
    Third,
}

assert_eq!(Test::VARIANT_COUNT, 3);
```

## Usage

If you're using Cargo, just add it to your Cargo.toml:

```toml
[dependencies]
variant_count = "1.1"
```

Use `VariantCount` in the `derive` enum attribute.
!*/

use proc_macro::TokenStream;
use syn::DeriveInput;
use quote::quote;

#[proc_macro_derive(VariantCount)]
pub fn variant_count(input: TokenStream) -> TokenStream {
    let ast: DeriveInput = syn::parse(input).unwrap();
    let (vis, ty, generics) = (&ast.vis, &ast.ident, &ast.generics);
    let variants = match &ast.data {
        syn::Data::Enum(e) => &e.variants,
        _ => panic!("VariantCount can only be derived for enums"),
    };
    let count = variants.len();
    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();

    let tokens = quote! {
        impl #impl_generics #ty #ty_generics
            #where_clause
        {
            #vis const VARIANT_COUNT: usize = #count;
        }
    };
    tokens.into()
}