name_variant/
lib.rs

1// Copyright (c) 2022, Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3//! Generates methods to print the name of the enum variant.
4//!
5//! # Example
6//!
7//! ```rust
8//! use name_variant::NamedVariant;
9//!
10//! # macro_rules! dont_test { () => {
11//! #[derive(NamedVariant)]
12//! enum TestEnum {
13//!     A,
14//!     B(),
15//!     C(i32, i32),
16//!     D { _name: String, _age: i32 },
17//!     VariantTest,
18//! }
19//!
20//! let x = TestEnum::C(1, 2);
21//! assert_eq!(x.variant_name(), "C");
22//!
23//! let x = TestEnum::A;
24//! assert_eq!(x.variant_name(), "A");
25//!
26//! let x = TestEnum::B();
27//! assert_eq!(x.variant_name(), "B");
28//!
29//! let x = TestEnum::D {_name: "Jane Doe".into(), _age: 30 };
30//! assert_eq!(x.variant_name(), "D");
31//!
32//! let x = TestEnum::VariantTest;
33//! assert_eq!(x.variant_name(), "VariantTest");
34//!
35//! # }};
36//! ```
37
38extern crate proc_macro;
39
40use proc_macro::TokenStream;
41use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
42
43use quote::{quote, quote_spanned};
44use syn::spanned::Spanned;
45use syn::{parse_macro_input, Data, DataEnum, DeriveInput, Error, Fields};
46
47macro_rules! derive_error {
48    ($string: tt) => {
49        Error::new(Span::call_site(), $string)
50            .to_compile_error()
51            .into()
52    };
53}
54
55fn match_enum_to_string(name: &Ident, variants: &DataEnum) -> proc_macro2::TokenStream {
56    // the variant dispatch proper
57    let mut match_arms = quote! {};
58    for variant in variants.variants.iter() {
59        let variant_ident = &variant.ident;
60        let fields_in_variant = match &variant.fields {
61            Fields::Unnamed(_) => quote_spanned! {variant.span() => (..) },
62            Fields::Unit => quote_spanned! { variant.span() => },
63            Fields::Named(_) => quote_spanned! {variant.span() => {..} },
64        };
65        let variant_string = variant_ident.to_string();
66
67        match_arms.extend(quote! {
68            #name::#variant_ident #fields_in_variant => #variant_string,
69        });
70    }
71    match_arms
72}
73
74#[proc_macro_derive(NamedVariant)]
75pub fn derive_named_variant(input: TokenStream) -> TokenStream {
76    let input = parse_macro_input!(input as DeriveInput);
77
78    let name = &input.ident;
79    let data = &input.data;
80
81    let mut variant_checker_functions;
82
83    match data {
84        Data::Enum(data_enum) => {
85            variant_checker_functions = TokenStream2::new();
86
87            let variant_arms = match_enum_to_string(name, data_enum);
88
89            variant_checker_functions.extend(quote_spanned! { name.span() =>
90                const fn variant_name(&self) -> &'static str {
91                    match self {
92                        #variant_arms
93                    }
94                }
95            });
96        }
97        _ => return derive_error!("NamedVariant is only implemented for enums"),
98    };
99
100    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
101
102    let expanded = quote! {
103        impl #impl_generics #name #ty_generics #where_clause {
104            #variant_checker_functions
105        }
106    };
107
108    TokenStream::from(expanded)
109}
name_variant/lib.rs

name_variant/
lib.rs
