azalea_buf_macros/
write.rs

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

azalea_buf_macros/
write.rs
