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

azalea_buf_macros/
write.rs
