wasmtime_test_macros/
add_variants.rs

1use proc_macro2::{Span, TokenStream};
2use quote::quote;
3use syn::parse_macro_input;
4
5pub fn run(
6    attr: proc_macro::TokenStream,
7    item: proc_macro::TokenStream,
8) -> proc_macro::TokenStream {
9    expand_variants(
10        &parse_macro_input!(attr as syn::LitInt),
11        parse_macro_input!(item as syn::ItemEnum),
12    )
13    .unwrap_or_else(syn::Error::into_compile_error)
14    .into()
15}
16
17fn expand_variants(count: &syn::LitInt, mut ty: syn::ItemEnum) -> syn::Result<TokenStream> {
18    let count = count
19        .base10_digits()
20        .parse::<usize>()
21        .map_err(|_| syn::Error::new(count.span(), "expected unsigned integer"))?;
22
23    ty.variants = (0..count)
24        .map(|index| syn::Variant {
25            attrs: Vec::new(),
26            ident: syn::Ident::new(&format!("V{index}"), Span::call_site()),
27            fields: syn::Fields::Unit,
28            discriminant: None,
29        })
30        .collect();
31
32    Ok(quote!(#ty))
33}