educe/trait_handlers/deref_mut/
deref_mut_struct.rs

1use quote::quote;
2use syn::{spanned::Spanned, Data, DeriveInput, Field, Meta, Type};
3
4use super::{
5    models::{FieldAttributeBuilder, TypeAttributeBuilder},
6    TraitHandler,
7};
8use crate::{common::ident_index::IdentOrIndex, Trait};
9
10pub(crate) struct DerefMutStructHandler;
11
12impl TraitHandler for DerefMutStructHandler {
13    #[inline]
14    fn trait_meta_handler(
15        ast: &DeriveInput,
16        token_stream: &mut proc_macro2::TokenStream,
17        traits: &[Trait],
18        meta: &Meta,
19    ) -> syn::Result<()> {
20        let _ = TypeAttributeBuilder {
21            enable_flag: true
22        }
23        .build_from_deref_mut_meta(meta)?;
24
25        let mut deref_mut_token_stream = proc_macro2::TokenStream::new();
26
27        if let Data::Struct(data) = &ast.data {
28            let (index, field) = {
29                let fields = &data.fields;
30
31                if fields.len() == 1 {
32                    let field = fields.into_iter().next().unwrap();
33
34                    let _ = FieldAttributeBuilder {
35                        enable_flag: true
36                    }
37                    .build_from_attributes(&field.attrs, traits)?;
38
39                    (0usize, field)
40                } else {
41                    let mut deref_field: Option<(usize, &Field)> = None;
42
43                    for (index, field) in fields.iter().enumerate() {
44                        let field_attribute = FieldAttributeBuilder {
45                            enable_flag: true
46                        }
47                        .build_from_attributes(&field.attrs, traits)?;
48
49                        if field_attribute.flag {
50                            if deref_field.is_some() {
51                                return Err(super::panic::multiple_deref_mut_fields(
52                                    field_attribute.span,
53                                ));
54                            }
55
56                            deref_field = Some((index, field));
57                        }
58                    }
59
60                    if let Some(deref_field) = deref_field {
61                        deref_field
62                    } else {
63                        return Err(super::panic::no_deref_mut_field(meta.span()));
64                    }
65                }
66            };
67
68            let field_name = IdentOrIndex::from_ident_with_index(field.ident.as_ref(), index);
69
70            deref_mut_token_stream.extend(if let Type::Reference(_) = &field.ty {
71                quote! (self.#field_name)
72            } else {
73                quote! (&mut self.#field_name)
74            });
75        }
76
77        let ident = &ast.ident;
78
79        let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
80
81        token_stream.extend(quote! {
82            impl #impl_generics ::core::ops::DerefMut for #ident #ty_generics #where_clause {
83                #[inline]
84                fn deref_mut(&mut self) -> &mut Self::Target {
85                    #deref_mut_token_stream
86                }
87            }
88        });
89
90        Ok(())
91    }
92}