A procedural macro that could count the occurrences of a particular enum variant in a collection

Question

I'm trying to make a procedural macro (CountOf) that counts the ocurrences of a particular enum variant in a collection, in this case a vector.

This src/macros.rs is what I'm using for testing the macro

#[derive(count_of::CountOf, Clone)]
pub enum SomeEnum {
    Variant1,
    Variant2,
    Variant3,
}

#[cfg(test)]
mod tests {
    use super::*;
#[test]
    fn count_of_works() {
        use SomeEnum::*;
        let enums = vec![
            Variant1,
            Variant1,
            Variant2,
            Variant3,
            Variant2,
            Variant2,
        ];

        assert_eq!(enums.variant1_count(), 2);
        assert_eq!(enums.variant2_count(), 3);
        assert_eq!(enums.variant3_count(), 1);
    }
}

This is my lib.rs in the count_of crate:

use inflector::Inflector;
use quote::quote;

#[proc_macro_derive(CountOf)]
pub fn count_of(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input = syn::parse_macro_input!(input as syn::ItemEnum);
    let name = input.ident;

    let variants = input.variants.iter().map(|variant| {
        let variant_name = &variant.ident;
        let variant_count = variant_name.to_string().to_lowercase();

        quote! {
            pub fn #variant_count(&self) -> usize {
                self.iter().filter(|&&x| x == &#name::#variant_name).count()
            }
        }
    });

    let output = quote! {
        impl #name {
            #(#variants)*
        }
    };

    proc_macro::TokenStream::from(output)
}

This are my dependencies in my Cargo.toml:

[dependencies]
count-of           = { path = "./count-of" }

An this is the :

[package]
name    = "count-of"
version = "0.1.0"

[lib]
proc-macro = true

[dependencies]
syn         = { version = "2.0.15", features = ["full"] }
quote       = "1.0.26"
Inflector   = "0.11.4"
proc-macro2 = "1.0.56"

Getting this errors:

error: expected identifier, found `"Variant1"`
  --> src/macros.rs:92:10
   |
92 | #[derive(count_of::CountOf, Clone)]
   |          ^^^^^^^^^^^^^^^^^
   |          |
   |          expected identifier
   |          while parsing this item list starting here
   |          the item list ends here
   |
   = note: this error originates in the derive macro `count_of::CountOf` (in Nightly builds, run with -Z macro-backtrace for more info)

error: proc-macro derive produced unparseable tokens
  --> src/macros.rs:92:10
   |
92 | #[derive(count_of::CountOf, Clone)]
   |          ^^^^^^^^^^^^^^^^^
error[E0599]: no method named `variant1_count` found for struct `Vec<SomeEnum>` in the current scope
   --> src/macros.rs:135:23
    |
135 |         assert_eq!(enums.variant1_count(), 2);
    |                             ^^^^^^^^^^^^^^^^^^^^^^^^^^ method not found in `Vec<SomeEnum>`

error[E0599]: no method named `variant2_count` found for struct `Vec<SomeEnum>` in the current scope
   --> src/macros.rs:136:23
    |
136 |         assert_eq!(enums.variant2_count(), 3);
    |                             ^^^^^^^^^^^^^^^^^ method not found in `Vec<SomeEnum>`

error[E0599]: no method named `variant3_count` found for struct `Vec<SomeEnum>` in the current scope
   --> src/macros.rs:137:23
    |
137 |         assert_eq!(enums.variant3_count(), 1);
    |                             ^^^^^^^^ method not found in `Vec<SomeEnum>`

I'm thinking my procedural macro is the issue but can't find how to fix it. If someone has a solution would really appreciate it.

Answer 1

There are a couple of errors in your code.

• Don't use Strings inside of quote!. Use the proper type instead, in this case Ident.
• Don't implement the functions for the type itself, if you want to call them on a Vec.

The second one is the hard one - because you cannot implement methods on a foreign type.

The solution is to create a new trait and then implement the trait for the foreign type, here Vec<SomeEnum>.

Here's some working code:

use proc_macro2::Ident;
use quote::{format_ident, quote};

#[proc_macro_derive(CountOf)]
pub fn count_of(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input = syn::parse_macro_input!(input as syn::ItemEnum);
    let name = input.ident;

    let trait_name = format_ident!("{}VecExt", name);

    let variants = input.variants.iter().map(|variant| {
        let variant_name = &variant.ident;
        let variant_count = variant_name.to_string().to_lowercase() + "_count";
        let variant_count_ident = Ident::new(&variant_count, variant_name.span());

        quote! {
            fn #variant_count_ident(&self) -> usize {
                self.as_ref().iter().filter(|&x| x == &#name::#variant_name).count()
            }
        }
    });

    let output = quote! {
        pub trait #trait_name: AsRef<[#name]> {
            #(#variants)*
        }
        impl<T> #trait_name for T where T: AsRef<[#name]> {}
    };

    proc_macro::TokenStream::from(output)
}

use rust_playground::CountOf;

#[derive(CountOf, Clone, Eq, PartialEq)]
pub enum SomeEnum {
    Variant1,
    Variant2,
    Variant3,
}

fn main() {
    use SomeEnum::*;

    let enums = vec![Variant1, Variant1, Variant2, Variant3, Variant2, Variant2];

    println!("Variant 1: {}", enums.variant1_count());
    println!("Variant 2: {}", enums.variant2_count());
    println!("Variant 3: {}", enums.variant3_count());
}

Variant 1: 2
Variant 2: 3
Variant 3: 1

When putting the code through cargo expand, you can see what it exands to:

pub enum SomeEnum {
    Variant1,
    Variant2,
    Variant3,
}
pub trait SomeEnumVecExt: AsRef<[SomeEnum]> {
    fn variant1_count(&self) -> usize {
        self.as_ref().iter().filter(|&x| x == &SomeEnum::Variant1).count()
    }
    fn variant2_count(&self) -> usize {
        self.as_ref().iter().filter(|&x| x == &SomeEnum::Variant2).count()
    }
    fn variant3_count(&self) -> usize {
        self.as_ref().iter().filter(|&x| x == &SomeEnum::Variant3).count()
    }
}
impl<T> SomeEnumVecExt for T
where
    T: AsRef<[SomeEnum]>,
{}

Also note the #[derive(Eq, PartialEq)], otherwise you cannot compare enums.

Maybe a little explanation:

• I create the trait MyEnumVecExt
• I implement the trait for all types that implement AsRef([SomeEnum]), which is all the types that can be referenced as &[SomeEnum], like Vec<SomeEnum>, [SomeEnum; 10], &mut [SomeEnum], etc.