educe/trait_handlers/deref_mut/
deref_mut_enum.rs

1use std::{fmt::Write, str::FromStr};
2
3use proc_macro2::TokenStream;
4use quote::quote;
5use syn::{Data, DeriveInput, Fields, Meta};
6
7use super::{
8    super::TraitHandler,
9    models::{FieldAttributeBuilder, TypeAttributeBuilder},
10};
11use crate::{panic, Trait};
12
13pub struct DerefMutEnumHandler;
14
15impl TraitHandler for DerefMutEnumHandler {
16    fn trait_meta_handler(
17        ast: &DeriveInput,
18        tokens: &mut TokenStream,
19        traits: &[Trait],
20        meta: &Meta,
21    ) {
22        let _ = TypeAttributeBuilder {
23            enable_flag: true
24        }
25        .from_deref_mut_meta(meta);
26
27        let enum_name = ast.ident.to_string();
28
29        let mut deref_mut_tokens = TokenStream::new();
30
31        let mut match_tokens = String::from("match self {");
32
33        if let Data::Enum(data) = &ast.data {
34            for variant in data.variants.iter() {
35                let _ = TypeAttributeBuilder {
36                    enable_flag: false
37                }
38                .from_attributes(&variant.attrs, traits);
39
40                let variant_ident = variant.ident.to_string();
41
42                match &variant.fields {
43                    Fields::Unit => {
44                        // TODO Unit
45                        panic::deref_mut_cannot_support_unit_variant();
46                    },
47                    Fields::Named(fields) => {
48                        // TODO Struct
49                        let mut pattern_tokens = String::new();
50                        let mut block_tokens = String::new();
51
52                        let mut counter = 0;
53
54                        for field in fields.named.iter() {
55                            let field_attribute = FieldAttributeBuilder {
56                                enable_flag: true
57                            }
58                            .from_attributes(&field.attrs, traits);
59
60                            if field_attribute.flag {
61                                if !block_tokens.is_empty() {
62                                    panic::multiple_deref_mut_fields_of_variant(&variant_ident);
63                                }
64
65                                let field_name = field.ident.as_ref().unwrap().to_string();
66
67                                block_tokens
68                                    .write_fmt(format_args!(
69                                        "return {field_name};",
70                                        field_name = field_name
71                                    ))
72                                    .unwrap();
73                                pattern_tokens
74                                    .write_fmt(format_args!(
75                                        "{field_name}, ..",
76                                        field_name = field_name
77                                    ))
78                                    .unwrap();
79                            }
80
81                            counter += 1;
82                        }
83
84                        if block_tokens.is_empty() {
85                            if counter == 1 {
86                                let field = fields.named.iter().next().unwrap();
87
88                                let field_name = field.ident.as_ref().unwrap().to_string();
89
90                                block_tokens
91                                    .write_fmt(format_args!(
92                                        "return {field_name};",
93                                        field_name = field_name
94                                    ))
95                                    .unwrap();
96                                pattern_tokens
97                                    .write_fmt(format_args!(
98                                        "{field_name}, ..",
99                                        field_name = field_name
100                                    ))
101                                    .unwrap();
102                            } else {
103                                panic::no_deref_mut_field_of_variant(&variant_ident);
104                            }
105                        }
106
107                        match_tokens
108                            .write_fmt(format_args!(
109                                "{enum_name}::{variant_ident} {{ {pattern_tokens} }} => {{ \
110                                 {block_tokens} }}",
111                                enum_name = enum_name,
112                                variant_ident = variant_ident,
113                                pattern_tokens = pattern_tokens,
114                                block_tokens = block_tokens
115                            ))
116                            .unwrap();
117                    },
118                    Fields::Unnamed(fields) => {
119                        // TODO Tuple
120                        let mut pattern_tokens = String::new();
121                        let mut block_tokens = String::new();
122
123                        let mut counter = 0;
124
125                        for (index, field) in fields.unnamed.iter().enumerate() {
126                            let field_attribute = FieldAttributeBuilder {
127                                enable_flag: true
128                            }
129                            .from_attributes(&field.attrs, traits);
130
131                            if field_attribute.flag {
132                                if !block_tokens.is_empty() {
133                                    panic::multiple_deref_mut_fields_of_variant(&variant_ident);
134                                }
135
136                                let field_name = format!("{}", index);
137
138                                block_tokens
139                                    .write_fmt(format_args!(
140                                        "return _{field_name};",
141                                        field_name = field_name
142                                    ))
143                                    .unwrap();
144                                pattern_tokens
145                                    .write_fmt(format_args!(
146                                        "_{field_name},",
147                                        field_name = field_name
148                                    ))
149                                    .unwrap();
150                            } else {
151                                pattern_tokens.push_str("_,");
152                            }
153
154                            counter += 1;
155                        }
156
157                        if block_tokens.is_empty() {
158                            if counter == 1 {
159                                let field_name = String::from("0");
160
161                                block_tokens
162                                    .write_fmt(format_args!(
163                                        "return _{field_name};",
164                                        field_name = field_name
165                                    ))
166                                    .unwrap();
167
168                                pattern_tokens.clear();
169                                pattern_tokens
170                                    .write_fmt(format_args!(
171                                        "_{field_name}",
172                                        field_name = field_name
173                                    ))
174                                    .unwrap();
175                            } else {
176                                panic::no_deref_mut_field_of_variant(&variant_ident);
177                            }
178                        }
179
180                        match_tokens
181                            .write_fmt(format_args!(
182                                "{enum_name}::{variant_ident}( {pattern_tokens} ) => {{ \
183                                 {block_tokens} }}",
184                                enum_name = enum_name,
185                                variant_ident = variant_ident,
186                                pattern_tokens = pattern_tokens,
187                                block_tokens = block_tokens
188                            ))
189                            .unwrap();
190                    },
191                }
192            }
193        }
194
195        match_tokens.push('}');
196
197        deref_mut_tokens.extend(TokenStream::from_str(&match_tokens).unwrap());
198
199        let ident = &ast.ident;
200
201        let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
202
203        let deref_mut_impl = quote! {
204            impl #impl_generics core::ops::DerefMut for #ident #ty_generics #where_clause {
205                #[inline]
206                fn deref_mut(&mut self) -> &mut Self::Target {
207                    #deref_mut_tokens
208                }
209            }
210        };
211
212        tokens.extend(deref_mut_impl);
213    }
214}
educe/trait_handlers/deref_mut/deref_mut_enum.rs

educe/trait_handlers/deref_mut/
deref_mut_enum.rs
