risingwave_common_proc_macro/
lib.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#![feature(coverage_attribute)]
16
17use estimate_size::{
18    add_trait_bounds, extract_ignored_generics_list, has_nested_flag_attribute_list,
19};
20use proc_macro::TokenStream;
21use proc_macro_error::proc_macro_error;
22use quote::quote;
23use syn::parse_macro_input;
24
25mod config;
26mod config_doc;
27mod estimate_size;
28mod session_config;
29
30/// Sections in the configuration file can use `#[derive(OverrideConfig)]` to generate the
31/// implementation of overwriting configs from the file.
32///
33/// In the struct definition, use `#[override_opts(path = ...)]` on a field to indicate the field in
34/// `RwConfig` to override.
35///
36/// An example:
37///
38/// ```ignore
39/// #[derive(OverrideConfig)]
40/// struct Opts {
41///     #[override_opts(path = meta.listen_addr)]
42///     listen_addr: Option<String>,
43/// }
44/// ```
45///
46/// will generate
47///
48/// ```ignore
49/// impl OverrideConfig for Opts {
50///     fn r#override(self, config: &mut RwConfig) {
51///         if let Some(v) = self.required_str {
52///             config.meta.listen_addr = v;
53///         }
54///     }
55/// }
56/// ```
57#[proc_macro_derive(OverrideConfig, attributes(override_opts))]
58#[proc_macro_error]
59pub fn override_config(input: TokenStream) -> TokenStream {
60    let input = parse_macro_input!(input);
61
62    let r#gen = config::produce_override_config(input);
63
64    r#gen.into()
65}
66
67/// `EstimateSize` can be derived if when all the fields in a
68/// struct or enum can implemented `EstimateSize`.
69#[proc_macro_derive(EstimateSize, attributes(estimate_size))]
70pub fn derive_estimate_size(input: TokenStream) -> TokenStream {
71    // Construct a representation of Rust code as a syntax tree
72    // that we can manipulate
73    let ast: syn::DeriveInput = syn::parse(input).unwrap();
74
75    // The name of the struct.
76    let name = &ast.ident;
77
78    // Extract all generics we shall ignore.
79    let ignored = extract_ignored_generics_list(&ast.attrs);
80
81    // Add a bound `T: EstimateSize` to every type parameter T.
82    let generics = add_trait_bounds(ast.generics, &ignored);
83
84    // Extract the generics of the struct/enum.
85    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
86
87    // Traverse the parsed data to generate the individual parts of the function.
88    match ast.data {
89        syn::Data::Enum(data_enum) => {
90            if data_enum.variants.is_empty() {
91                // Empty enums are easy to implement.
92                let r#gen = quote! {
93                    impl EstimateSize for #name {
94                        fn estimated_heap_size(&self) -> usize {
95                            0
96                        }
97                    }
98                };
99                return r#gen.into();
100            }
101
102            let mut cmds = Vec::with_capacity(data_enum.variants.len());
103
104            for variant in &data_enum.variants {
105                let ident = &variant.ident;
106
107                match &variant.fields {
108                    syn::Fields::Unnamed(unnamed_fields) => {
109                        let num_fields = unnamed_fields.unnamed.len();
110
111                        let mut field_idents = Vec::with_capacity(num_fields);
112                        for i in 0..num_fields {
113                            let field_ident = String::from("v") + &i.to_string();
114                            let field_ident = syn::parse_str::<syn::Ident>(&field_ident).unwrap();
115
116                            field_idents.push(field_ident);
117                        }
118
119                        let mut field_cmds = Vec::with_capacity(num_fields);
120
121                        for (i, _field) in unnamed_fields.unnamed.iter().enumerate() {
122                            let field_ident = String::from("v") + &i.to_string();
123                            let field_ident = syn::parse_str::<syn::Ident>(&field_ident).unwrap();
124
125                            field_cmds.push(quote! {
126                                total += EstimateSize::estimated_heap_size(#field_ident);
127                            })
128                        }
129
130                        cmds.push(quote! {
131                            Self::#ident(#(#field_idents,)*) => {
132                                let mut total = 0;
133
134                                #(#field_cmds)*;
135
136                                total
137                            }
138                        });
139                    }
140                    syn::Fields::Named(named_fields) => {
141                        let num_fields = named_fields.named.len();
142
143                        let mut field_idents = Vec::with_capacity(num_fields);
144
145                        let mut field_cmds = Vec::with_capacity(num_fields);
146
147                        for field in &named_fields.named {
148                            let field_ident = field.ident.as_ref().unwrap();
149
150                            field_idents.push(field_ident);
151
152                            field_cmds.push(quote! {
153                                total += #field_ident.estimated_heap_size();
154                            })
155                        }
156
157                        cmds.push(quote! {
158                            Self::#ident{#(#field_idents,)*} => {
159                                let mut total = 0;
160
161                                #(#field_cmds)*;
162
163                                total
164                            }
165                        });
166                    }
167                    syn::Fields::Unit => {
168                        cmds.push(quote! {
169                            Self::#ident => 0,
170                        });
171                    }
172                }
173            }
174
175            // Build the trait implementation
176            let r#gen = quote! {
177                impl #impl_generics EstimateSize for #name #ty_generics #where_clause {
178                    fn estimated_heap_size(&self) -> usize {
179                        match self {
180                            #(#cmds)*
181                        }
182                    }
183                }
184            };
185            r#gen.into()
186        }
187        syn::Data::Union(_data_union) => {
188            panic!("Deriving EstimateSize for unions is currently not supported.")
189        }
190        syn::Data::Struct(data_struct) => {
191            if data_struct.fields.is_empty() {
192                // Empty structs are easy to implement.
193                let r#gen = quote! {
194                    impl EstimateSize for #name {
195                        fn estimated_heap_size(&self) -> usize {
196                            0
197                        }
198                    }
199                };
200                return r#gen.into();
201            }
202
203            let mut field_cmds = Vec::with_capacity(data_struct.fields.len());
204
205            match data_struct.fields {
206                syn::Fields::Unnamed(unnamed_fields) => {
207                    for (i, field) in unnamed_fields.unnamed.iter().enumerate() {
208                        // Check if the value should be ignored. If so skip it.
209                        if has_nested_flag_attribute_list(&field.attrs, "estimate_size", "ignore") {
210                            continue;
211                        }
212
213                        let idx = syn::Index::from(i);
214                        field_cmds.push(quote! {
215                            total += EstimateSize::estimated_heap_size(&self.#idx);
216                        })
217                    }
218                }
219                syn::Fields::Named(named_fields) => {
220                    for field in &named_fields.named {
221                        // Check if the value should be ignored. If so skip it.
222                        if has_nested_flag_attribute_list(&field.attrs, "estimate_size", "ignore") {
223                            continue;
224                        }
225
226                        let field_ident = field.ident.as_ref().unwrap();
227                        field_cmds.push(quote! {
228                            total += &self.#field_ident.estimated_heap_size();
229                        })
230                    }
231                }
232                syn::Fields::Unit => {}
233            }
234
235            // Build the trait implementation
236            let r#gen = quote! {
237                impl #impl_generics EstimateSize for #name #ty_generics #where_clause {
238                    fn estimated_heap_size(&self) -> usize {
239                        let mut total = 0;
240
241                        #(#field_cmds)*;
242
243                        total
244                    }
245                }
246            };
247            r#gen.into()
248        }
249    }
250}
251
252/// To add a new parameter, you can add a field with `#[parameter]` in the struct
253/// A default value is required by setting the `default` option.
254/// The field name will be the parameter name. You can overwrite the parameter name by setting the `rename` option.
255/// To check the input parameter, you can use `check_hook` option.
256///
257/// `flags` options include
258/// - `SETTER`: to manually write a `set_your_parameter_name` function, in which you should call `set_your_parameter_name_inner`.
259/// - `REPORT`: to report the parameter through `ConfigReporter`
260/// - `NO_ALTER_SYS`: disallow the parameter to be set by `alter system set`
261#[proc_macro_derive(SessionConfig, attributes(parameter))]
262#[proc_macro_error]
263pub fn session_config(input: TokenStream) -> TokenStream {
264    let input = parse_macro_input!(input);
265    session_config::derive_config(input).into()
266}
267
268/// This proc macro recursively extracts rustdoc comments from the fields in a struct and generates a method
269/// that produces docs for each field.
270///
271/// Unlike rustdoc, this tool focuses solely on extracting rustdoc for struct fields, without methods.
272///
273/// Example:
274///
275/// ```ignore
276/// #[derive(ConfigDoc)]
277/// pub struct Foo {
278///   /// Description for `a`.
279///   a: i32,
280///
281///   #[config_doc(nested)]
282///   b: Bar,
283///
284///   #[config_doc(omitted)]
285///   dummy: (),
286/// }
287/// ```
288///
289/// The `#[config_doc(nested)]` attribute indicates that the field is a nested config that will be documented in a separate section.
290/// Fields marked with `#[config_doc(omitted)]` will simply be omitted from the doc.
291///
292/// Here is the method generated by this macro:
293///
294/// ```ignore
295/// impl Foo {
296///     pub fn config_docs(name: String, docs: &mut std::collections::BTreeMap<String, Vec<(String, String)>>)
297/// }
298/// ```
299///
300/// In `test_example_up_to_date`, we further process the output of this method to generate a markdown in src/config/docs.md.
301#[proc_macro_derive(ConfigDoc, attributes(config_doc))]
302pub fn config_doc(input: TokenStream) -> TokenStream {
303    let input = parse_macro_input!(input);
304
305    let r#gen = config_doc::generate_config_doc_fn(input);
306
307    r#gen.into()
308}