How to get all the variants of an enum in a Vec<T> with a proc-macro?

user_14075508 0
Viewed 1067

I am wondering how to implement a method for any enum that will return all variants in a Vec<T> or some kind of collection type. Something like:

pub enum MyEnum {
    EnumVariant1
    EnumVariant2
    ...
}

impl MyEnum {
    fn values(&self) -> Vec<MyEnum> {
        // do Rust stuff here
    }
}
rust enums
2 Answers

There's no smart and standard solution.

An obvious one is to declare the array yourself, by repeating the variants:

static VARIANTS: &[MyEnum] = &[
    MyEnum::EnumVariant1,
    MyEnum::EnumVariant2,
];

This is a reasonable solution. When your code evolve, it's frequent you discover you don't want all variants in your static array. Or that you need several arrays.

Alternatively, if there are many elements or several enums, you may create a macro for this purpose:

macro_rules! make_enum {
    (
        $name:ident $array:ident {
            $( $variant:ident, )*
        }
    ) => {
        pub enum $name {
            $( $variant, )*
        }
        static $array: &[$name] = &[
            $( $name::$variant, )*
        ];
    }
}

make_enum! (MyEnum VARIANTS {
    EnumVariant1,
    EnumVariant2,
});

This make_enum! macro call would create both the enum and a static array called VARIANTS.

user_2635250

Alternatively you can use a proc macro which makes it generic (inspired by How to get the number of elements (variants) in an enum as a constant value?):

extern crate proc_macro;
extern crate syn;
#[macro_use]
extern crate quote;

use proc_macro::TokenStream;

#[proc_macro_derive(AllVariants)]
pub fn derive_all_variants(input: TokenStream) -> TokenStream {
    let syn_item: syn::DeriveInput = syn::parse(input).unwrap();

    let variants = match syn_item.data {
        syn::Data::Enum(enum_item) => {
            enum_item.variants.into_iter().map(|v| v.ident)
        }
        _ => panic!("AllVariants only works on enums"),
    };
    let enum_name = syn_item.ident;

    let expanded = quote! {
        impl #enum_name {
            pub fn all_variants() -> &'static[#enum_name] {
                &[ #(#enum_name::#variants),* ]
            }
        }
    };
    expanded.into()
}

Then you can use it like this:

use all_variants::AllVariants;

#[derive(AllVariants, Debug)]
enum Direction {
    Left,
    Top,
    Right,
    Bottom,
}

fn main() {
    println!("{:?}", Direction::all_variants());
}

Output:

[Left, Top, Right, Bottom]

But as @Denys Séguret wrote in a previous comment, it's lot more work:

  • proc macros has to reside in their own crate (if you want to publish your work on crates.io, the proc macro crate has to be published as well);
  • it requires syn and quote crates in the proc macro (if don't want to reinvent the wheel).
user_32716870
Related