azalea_buf_macros/
write.rs

1use proc_macro2::Span;
2use quote::{ToTokens, quote};
3use syn::{Data, Field, FieldsNamed, Generics, Ident, punctuated::Punctuated, token::Comma};
4
5pub fn create_impl_azaleawrite(
6    ident: &Ident,
7    generics: &Generics,
8    data: &Data,
9) -> proc_macro2::TokenStream {
10    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
11
12    match data {
13        syn::Data::Struct(syn::DataStruct { fields, .. }) => match fields {
14            syn::Fields::Named(FieldsNamed { named, .. }) => {
15                let write_fields =
16                    write_named_fields(named, Some(&Ident::new("self", Span::call_site())));
17
18                quote! {
19                    impl #impl_generics azalea_buf::AzaleaWrite for #ident #ty_generics #where_clause {
20                        fn azalea_write(&self, buf: &mut impl std::io::Write) -> std::result::Result<(), std::io::Error> {
21                            #write_fields
22                            Ok(())
23                        }
24                    }
25                }
26            }
27            syn::Fields::Unit => {
28                quote! {
29                    impl #impl_generics azalea_buf::AzaleaWrite for #ident #ty_generics #where_clause {
30                        fn azalea_write(&self, buf: &mut impl std::io::Write) -> std::result::Result<(), std::io::Error> {
31                            Ok(())
32                        }
33                    }
34                }
35            }
36            syn::Fields::Unnamed(fields) => {
37                let write_fields = write_unnamed_fields(&fields.unnamed);
38
39                quote! {
40                    impl #impl_generics azalea_buf::AzaleaWrite for #ident #ty_generics #where_clause {
41                        fn azalea_write(&self, buf: &mut impl std::io::Write) -> std::result::Result<(), std::io::Error> {
42                            #write_fields
43                            Ok(())
44                        }
45                    }
46                }
47            }
48        },
49        syn::Data::Enum(syn::DataEnum { variants, .. }) => {
50            // remember whether it's a data variant so we can do an optimization later
51            let mut is_data_enum = false;
52            let mut match_arms = quote!();
53            let mut match_arms_without_id = quote!();
54            let mut variant_discrim: u32 = 0;
55            let mut first = true;
56            for variant in variants {
57                match &variant.discriminant.as_ref() {
58                    Some(d) => {
59                        variant_discrim = match &d.1 {
60                            syn::Expr::Lit(e) => match &e.lit {
61                                syn::Lit::Int(i) => i.base10_parse().unwrap(),
62                                // syn::Lit::Str(s) => s.value(),
63                                _ => panic!("Error parsing enum discriminant as int (is {e:?})"),
64                            },
65                            syn::Expr::Unary(_) => {
66                                panic!("Negative enum discriminants are not supported")
67                            }
68                            _ => {
69                                panic!("Error parsing enum discriminant as literal (is {:?})", d.1)
70                            }
71                        }
72                    }
73                    None => {
74                        if first {
75                            first = false;
76                        } else {
77                            variant_discrim += 1;
78                        }
79                    }
80                }
81
82                let variant_name = &variant.ident;
83
84                // the variant number that we're going to write
85                let write_the_variant = quote! {
86                    azalea_buf::AzaleaWriteVar::azalea_write_var(&#variant_discrim, buf)?;
87                };
88                match &variant.fields {
89                    syn::Fields::Named(f) => {
90                        is_data_enum = true;
91                        let field_names = f
92                            .named
93                            .iter()
94                            .map(|f| f.ident.clone().unwrap())
95                            .collect::<Vec<_>>();
96                        let write_fields = write_named_fields(&f.named, None);
97                        match_arms.extend(quote! {
98                            Self::#variant_name { #(#field_names),* } => {
99                                #write_the_variant
100                                #write_fields
101                            }
102                        });
103                        match_arms_without_id.extend(quote! {
104                            Self::#variant_name { #(#field_names),* } => {
105                                #write_fields
106                            }
107                        });
108                    }
109                    syn::Fields::Unit => {
110                        match_arms.extend(quote! {
111                            Self::#variant_name => {
112                                #write_the_variant
113                            }
114                        });
115                        match_arms_without_id.extend(quote! {
116                            Self::#variant_name => {}
117                        });
118                    }
119                    syn::Fields::Unnamed(fields) => {
120                        is_data_enum = true;
121                        let mut writers_code = quote! {};
122                        let mut params_code = quote! {};
123                        for (i, f) in fields.unnamed.iter().enumerate() {
124                            let param_ident = Ident::new(&format!("data{i}"), Span::call_site());
125                            params_code.extend(quote! { #param_ident, });
126                            if f.attrs.iter().any(|attr| attr.path().is_ident("var")) {
127                                writers_code.extend(quote! {
128                                    azalea_buf::AzaleaWriteVar::azalea_write_var(#param_ident, buf)?;
129                                });
130                            } else {
131                                writers_code.extend(quote! {
132                                    azalea_buf::AzaleaWrite::azalea_write(#param_ident, buf)?;
133                                });
134                            }
135                        }
136                        match_arms.extend(quote! {
137                            Self::#variant_name(#params_code) => {
138                                #write_the_variant
139                                #writers_code
140                            }
141                        });
142                        match_arms_without_id.extend(quote! {
143                            Self::#variant_name(data) => {
144                                azalea_buf::AzaleaWrite::azalea_write(data, buf)?;
145                            }
146                        });
147                    }
148                }
149            }
150            if is_data_enum {
151                quote! {
152                    impl #impl_generics azalea_buf::AzaleaWrite for #ident #ty_generics #where_clause {
153                        fn azalea_write(&self, buf: &mut impl std::io::Write) -> std::result::Result<(), std::io::Error> {
154                            match self {
155                                #match_arms
156                            }
157                            Ok(())
158                        }
159                    }
160                    impl #impl_generics #ident #ty_generics #where_clause {
161                        pub fn write_without_id(&self, buf: &mut impl std::io::Write) -> std::result::Result<(), std::io::Error> {
162                            match self {
163                                #match_arms_without_id
164                            }
165                            Ok(())
166                        }
167                    }
168                }
169            } else {
170                // optimization: if it doesn't have data we can just do `as u32`
171                quote! {
172                    impl #impl_generics azalea_buf::AzaleaWrite for #ident #ty_generics #where_clause {
173                        fn azalea_write(&self, buf: &mut impl std::io::Write) -> std::result::Result<(), std::io::Error> {
174                            azalea_buf::AzaleaWriteVar::azalea_write_var(&(*self as u32), buf)
175                        }
176                    }
177                }
178            }
179        }
180        _ => panic!("#[derive(AzBuf)] can only be used on structs"),
181    }
182}
183
184fn write_named_fields(
185    named: &Punctuated<Field, Comma>,
186    ident_name: Option<&Ident>,
187) -> proc_macro2::TokenStream {
188    let write_fields = named.iter().map(|f| {
189        let field_name = &f.ident;
190        let ident_dot_field = match ident_name {
191            Some(ident) => quote! { &#ident.#field_name },
192            None => quote! { #field_name },
193        };
194
195        make_write_call(f, ident_dot_field)
196    });
197    quote! { #(#write_fields)* }
198}
199
200fn write_unnamed_fields(named: &Punctuated<Field, Comma>) -> proc_macro2::TokenStream {
201    let write_fields = named.iter().enumerate().map(|(i, f)| {
202        let i_literal = syn::Index::from(i);
203        let ident_dot_field = quote! { &self.#i_literal };
204
205        make_write_call(f, ident_dot_field)
206    });
207    quote! { #(#write_fields)* }
208}
209
210fn make_write_call(
211    f: &Field,
212    ident_dot_field: proc_macro2::TokenStream,
213) -> proc_macro2::TokenStream {
214    let field_type = &f.ty;
215    // do a different buf.write_* for each field depending on the type
216    // if it's a string, use buf.write_string
217    match field_type {
218        syn::Type::Path(_) | syn::Type::Array(_) => {
219            if f.attrs.iter().any(|attr| attr.path().is_ident("var")) {
220                quote! {
221                    azalea_buf::AzaleaWriteVar::azalea_write_var(#ident_dot_field, buf)?;
222                }
223            } else {
224                quote! {
225                    azalea_buf::AzaleaWrite::azalea_write(#ident_dot_field, buf)?;
226                }
227            }
228        }
229        _ => panic!(
230            "Error writing field {:?}: {}",
231            f.ident,
232            field_type.to_token_stream()
233        ),
234    }
235}
azalea_buf_macros/write.rs

azalea_buf_macros/
write.rs
