enum-filter 0.1.0

Filter to a specific enum variant from an iterator
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

use quote::{format_ident, quote};
use syn::{parse_quote, DataEnum, Expr, Fields, Ident, ItemImpl, Pat, Type, Variant};

use crate::names::fn_name;
use crate::{functions, IterType};

fn get_filter_return_type(variant: &Variant, iter_type: IterType) -> Type {
    match &variant.fields {
        Fields::Unit => parse_quote! {()},
        Fields::Unnamed(fields) => {
            if fields.unnamed.len() == 1 {
                let t = &fields.unnamed[0].ty;
                match iter_type {
                    IterType::Owned => t.clone(),
                    IterType::Ref => parse_quote! {&'eflt #t},
                    IterType::RefMut => parse_quote! {&'eflt mut #t},
                }
            } else {
                let types = fields.unnamed.iter().map(|f| {
                    let t = &f.ty;

                    match iter_type {
                        IterType::Owned => t.clone(),
                        IterType::Ref => parse_quote! {&'eflt #t},
                        IterType::RefMut => parse_quote! {&'eflt mut #t},
                    }
                });
                parse_quote! {(#(#types,)*)}
            }
        }
        Fields::Named(fields) => {
            let ident = format_ident!("{}Data", variant.ident);
            let field_types = fields.named.iter().map(|f| &f.ty);

            match iter_type {
                IterType::Owned => parse_quote! {#ident<#(#field_types),*>},
                IterType::Ref => parse_quote! {#ident<#(&'eflt #field_types),*>},
                IterType::RefMut => parse_quote! {#ident<#(&'eflt mut #field_types),*>},
            }
        }
    }
}

fn get_filter_return_value(variant: &Variant) -> Expr {
    match &variant.fields {
        Fields::Unit => parse_quote!(()),

        Fields::Unnamed(fields) => {
            if fields.unnamed.len() == 1 {
                parse_quote! {v0}
            } else {
                let names = (0..fields.unnamed.len()).map(|i| format_ident!("v{}", i));
                parse_quote! { (#(#names,)*) }
            }
        }

        Fields::Named(fields) => {
            let names = fields.named.iter().map(|f| &f.ident);
            let name = format_ident!("{}Data", variant.ident);
            parse_quote! { #name { #(#names,)* } }
        }
    }
}

fn get_variant_matcher(enum_name: &Ident, variant: &Variant) -> Pat {
    let variant_name = &variant.ident;

    match &variant.fields {
        Fields::Unit => parse_quote! {#enum_name :: #variant_name},

        Fields::Unnamed(fields) => {
            let names = (0..fields.unnamed.len()).map(|i| format_ident!("v{}", i));
            parse_quote! { #enum_name :: #variant_name ( #(#names,)* ) }
        }

        Fields::Named(fields) => {
            let names = fields.named.iter().map(|f| &f.ident);
            parse_quote! { #enum_name :: #variant_name { #(#names,)* } }
        }
    }
}

pub(crate) fn generate_impl(
    enum_name: &Ident,
    trait_name: &Ident,
    enum_data: &DataEnum,
    iter_type: IterType,
) -> ItemImpl {
    let fns = enum_data.variants.iter().map(|v| {
        let fn_name = fn_name(&v.ident);
        let filter_arg: Type = match iter_type {
            IterType::Owned => parse_quote! {#enum_name},
            IterType::Ref => parse_quote! {&'eflt #enum_name},
            IterType::RefMut => parse_quote! {&'eflt mut #enum_name},
        };

        let filter_return = get_filter_return_type(v, iter_type);
        let matcher = get_variant_matcher(enum_name, v);
        let ret_value = get_filter_return_value(v);
        functions::generate_def(&fn_name, &filter_arg, &filter_return, &matcher, &ret_value)
    });

    let filter_return_types = enum_data
        .variants
        .iter()
        .map(|v| get_filter_return_type(v, iter_type));

    let iter_item: Type = match iter_type {
        IterType::Owned => parse_quote! {#enum_name},
        IterType::Ref => parse_quote! {&'eflt #enum_name},
        IterType::RefMut => parse_quote! {&'eflt mut #enum_name},
    };

    let impl_generics = match iter_type {
        IterType::Owned => quote! {<T: Iterator<Item = #iter_item>>},
        IterType::Ref | IterType::RefMut => quote! {<'eflt, T: Iterator<Item = #iter_item>>},
    };

    let trait_generics = quote! { <#iter_item, #(#filter_return_types),*>};

    parse_quote! {
        impl #impl_generics #trait_name #trait_generics for T {
            #(#fns)*
        }
    }
}