Skip to main content

enum_each_variant_derive/
lib.rs

1//! Derive method that returns each variant of an enum
2//!
3//! # Sample usage
4//!
5//! ```rust
6//! #[macro_use]
7//! extern crate enum_each_variant_derive;
8//!
9//! # fn main() {
10//! #[derive(EachVariant, Eq, PartialEq, Debug)]
11//! enum Thing {
12//!     One,
13//!     Two,
14//!     Three,
15//!     Four,
16//! }
17//!
18//! let all: Vec<Thing> = Thing::all_variants();
19//!
20//! assert_eq!(all, vec![Thing::One, Thing::Two, Thing::Three, Thing::Four]);
21//! # }
22//! ```
23//!
24//! # Gotcha
25//!
26//! Only works on enums where no variants have associated values. So we wouldn't be able to use it
27//! for this enum:
28//!
29//! ```rust
30//! # fn main() {
31//! enum TrainStatus {
32//!     OnTime,
33//!     DelayedBy(std::time::Duration),
34//! }
35//! # }
36//! ```
37
38extern crate proc_macro;
39extern crate syn;
40
41#[macro_use]
42extern crate quote;
43
44use proc_macro::TokenStream;
45use syn::*;
46
47#[doc(hidden)]
48#[proc_macro_derive(EachVariant)]
49pub fn each_variant(input: TokenStream) -> TokenStream {
50    let input: DeriveInput = parse(input).unwrap();
51    let expanded = impl_enum_each(input);
52    expanded.into()
53}
54
55fn impl_enum_each(ast: DeriveInput) -> quote::Tokens {
56    let name: &Ident = &ast.ident;
57
58    let enum_data: DataEnum = match ast.data {
59        Data::Enum(data) => data,
60        _ => panic!("#[derive(EachVariant)] is only defined for enums"),
61    };
62
63    let variants = enum_data.variants;
64    let variant_names = variants.iter().map(|ref variant| {
65        match variant.fields {
66            Fields::Unit => {}
67            _ => {
68                panic!("#[derive(EachVariant)] is only defined on enums where all the variants have no associated values");
69            }
70        };
71
72        variant.ident
73    });
74
75    let push_variants = variant_names
76        .map(|variant_name| {
77            quote! { acc.push(#name::#variant_name); }
78        });
79
80    quote! {
81        impl #name {
82            /// Build vector containing each variant of this enum
83            pub fn all_variants() -> Vec<Self> {
84                let mut acc: Vec<Self> = vec![];
85                #(#push_variants);*
86                acc
87            }
88        }
89    }
90}