enum_filter/
lib.rs

1//! `enum_filter` provides a macro that will generate "filter" methods for iterators over an enum
2//! It does so by creating a trait `EnumNameFilter` with the same visibilty as the enum
3//! For variants with named fields it will create a struct `VariantNameData`, also with the same visibility as the enum
4//!
5//! ```
6//! #[enum_filter]
7//! enum Example {
8//!     Variant1,
9//!     Variant2(u8),
10//!     Variant3 { data: u8 },
11//! }
12
13//! fn main() {
14//!     let mut test_vec = vec![
15//!         Example::Variant1,
16//!         Example::Variant2(2),
17//!         Example::Variant3 { data: 3 },
18//!     ];
19
20//!     assert_eq!(
21//!         test_vec.iter().filter_variant1().collect::<Vec<_>>(),
22//!         vec![()]
23//!     );
24
25//!     assert_eq!(
26//!         test_vec.iter_mut().filter_variant2().collect::<Vec<_>>(),
27//!         vec![&mut 2]
28//!     );
29
30//!     assert_eq!(
31//!         test_vec
32//!             .into_iter()
33//!             .filter_variant3()
34//!             .map(|v| v.data)
35//!             .collect::<Vec<_>>(),
36//!         vec![3]
37//!     );
38//! }
39//! ```
40
41use impls::generate_impl;
42use names::trait_name;
43use proc_macro::TokenStream;
44use quote::quote;
45use structs::generate_struct;
46use syn::{parse_macro_input, DeriveInput};
47use traits::generate_trait;
48
49mod functions;
50mod impls;
51mod names;
52mod structs;
53mod traits;
54
55#[derive(Clone, Copy)]
56enum IterType {
57    Owned,
58    Ref,
59    RefMut,
60}
61
62#[proc_macro_attribute]
63pub fn enum_filter(_args: TokenStream, input: TokenStream) -> TokenStream {
64    let ast = parse_macro_input!(input as DeriveInput);
65
66    let enum_data = match ast.data {
67        syn::Data::Enum(ref data) => data,
68        _ => panic!("enum_filter only supports enums"),
69    };
70
71    let enum_name = &ast.ident;
72
73    let structs = enum_data.variants.iter().filter_map(generate_struct);
74    let trait_name = trait_name(enum_name);
75    let trait_def = generate_trait(&ast.vis, &trait_name, enum_data);
76
77    let owned_impl = generate_impl(enum_name, &trait_name, enum_data, IterType::Owned);
78    let ref_impl = generate_impl(enum_name, &trait_name, enum_data, IterType::Ref);
79    let ref_mut_impl = generate_impl(enum_name, &trait_name, enum_data, IterType::RefMut);
80
81    quote! {
82        #ast
83
84        #(#structs)*
85
86        #trait_def
87
88        #owned_impl
89        #ref_impl
90        #ref_mut_impl
91    }
92    .into()
93}