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}