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

azalea_buf_macros/
write.rs
