sui_protocol_config_macros/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4extern crate proc_macro;
5
6use proc_macro::TokenStream;
7use quote::quote;
8use syn::{Data, DeriveInput, Fields, Type, parse_macro_input};
9
10/// This proc macro generates getters, attribute lookup, etc for protocol config fields of type `Option<T>`
11/// and for the feature flags
12/// Example for a field: `new_constant: Option<u64>`, and for feature flags `feature: bool`, we derive
13/// ```rust,ignore
14///     /// Returns the value of the field if exists at the given version, otherise panic
15///     pub fn new_constant(&self) -> u64 {
16///         self.new_constant.expect(Self::CONSTANT_ERR_MSG)
17///     }
18///     /// Returns the value of the field if exists at the given version, otherise None.
19///     pub fn new_constant_as_option(&self) -> Option<u64> {
20///         self.new_constant
21///     }
22///     // We auto derive an enum such that the variants are all the types of the fields
23///     pub enum ProtocolConfigValue {
24///        u32(u32),
25///        u64(u64),
26///        ..............
27///     }
28///     // This enum is used to return field values so that the type is also encoded in the response
29///
30///     /// Returns the value of the field if exists at the given version, otherise None
31///     pub fn lookup_attr(&self, value: String) -> Option<ProtocolConfigValue>;
32///
33///     /// Returns a map of all configs to values
34///     pub fn attr_map(&self) -> std::collections::BTreeMap<String, Option<ProtocolConfigValue>>;
35///
36///     /// Returns a feature by the string name or None if it doesn't exist
37///     pub fn lookup_feature(&self, value: String) -> Option<bool>;
38///
39///     /// Returns a map of all features to values
40///     pub fn feature_map(&self) -> std::collections::BTreeMap<String, bool>;
41/// ```
42///
43/// Every field (scalar and non-scalar) is also emitted into a typed
44/// `render<F: Format>(&self, meter: &mut impl Meter) -> Result<BTreeMap<String, F>, MeterError>`
45/// method, where each value is produced via `mysten_common::rpc_format::ToFormat`. Fields that
46/// aren't configured at the current protocol version render as `F::null(...)` rather than being
47/// absent from the map, so the keyset is stable across protocol versions. This is the path RPC
48/// code should use to expose protocol config to clients; the same call site can target
49/// `serde_json::Value`, `prost_types::Value`, or any other `Format` impl by choosing `F`.
50///
51/// Scalar (`u16`/`u32`/`u64`/`bool`) fields continue to feed `ProtocolConfigValue` / `attr_map`
52/// for back-compat with existing consumers. Non-scalar fields appear only in `render`. Add
53/// `#[skip_accessor]` to keep a field internal and out of every generated surface. Add
54/// `#[custom_setter]` to suppress the generated `set_x_for_testing` so the struct can provide a
55/// hand-written one (e.g. to validate the new value); the `from_str` setter and
56/// `set_attr_for_testing` still route through it by name.
57#[proc_macro_derive(ProtocolConfigAccessors, attributes(skip_accessor, custom_setter))]
58pub fn accessors_macro(input: TokenStream) -> TokenStream {
59    let ast = parse_macro_input!(input as DeriveInput);
60
61    let struct_name = &ast.ident;
62    let data = &ast.data;
63
64    let fields: Vec<AccessorField> = match data {
65        Data::Struct(data_struct) => match &data_struct.fields {
66            Fields::Named(fields_named) => fields_named
67                .named
68                .iter()
69                .filter_map(parse_accessor_field)
70                .collect(),
71            _ => panic!("Only named fields are supported."),
72        },
73        _ => panic!("Only structs supported."),
74    };
75
76    let expanded: Vec<ExpandedField> = fields.iter().map(expand_field).collect();
77
78    let accessors = expanded.iter().map(|e| &e.accessor);
79    let setters = expanded.iter().map(|e| &e.setter);
80    let render_arms = expanded.iter().map(|e| &e.render_arm);
81
82    // Scalar-only collections — driven by the optional `ScalarExtras` extension on each field.
83    let scalar_extras: Vec<&ScalarExtras> = expanded
84        .iter()
85        .filter_map(|e| e.scalar_extras.as_ref())
86        .collect();
87    let scalar_value_setter_arms = scalar_extras.iter().map(|s| &s.value_setter_arm);
88    let scalar_lookup_arms = scalar_extras.iter().map(|s| &s.lookup_arm);
89    let scalar_field_names = scalar_extras.iter().map(|s| &s.field_name_str);
90
91    // Multiple scalar fields of the same primitive type all share a single
92    // `ProtocolConfigValue` variant (e.g. every `u64` field maps to `ProtocolConfigValue::u64`).
93    let mut variant_decls = Vec::new();
94    let mut display_variants = Vec::new();
95    let mut seen = std::collections::HashSet::new();
96    for s in &scalar_extras {
97        if !seen.insert(s.variant_ident.to_string()) {
98            continue;
99        }
100        let ident = &s.variant_ident;
101        let inner = &s.inner_type;
102        variant_decls.push(quote! { #ident(#inner) });
103        display_variants.push(s.variant_ident.clone());
104    }
105
106    let output = quote! {
107        impl #struct_name {
108            const CONSTANT_ERR_MSG: &'static str = "protocol constant not present in current protocol version";
109            #(#accessors)*
110
111            /// Lookup a scalar config attribute by its string representation.
112            pub fn lookup_attr(&self, value: String) -> Option<ProtocolConfigValue> {
113                match value.as_str() {
114                    #(#scalar_lookup_arms)*
115                    _ => None,
116                }
117            }
118
119            /// Get a map of all scalar config attributes from string representations.
120            ///
121            /// Non-scalar (e.g. list-typed) fields aren't represented here — use
122            /// `Self::render` for a typed view that includes every field.
123            pub fn attr_map(&self) -> std::collections::BTreeMap<String, Option<ProtocolConfigValue>> {
124                vec![
125                    #(((#scalar_field_names).to_owned(), self.lookup_attr((#scalar_field_names).to_owned())),)*
126                    ].into_iter().collect()
127            }
128
129            /// Render every protocol-config attribute into the chosen `Format`.
130            ///
131            /// Fields that aren't configured at this protocol version render as `F::null(...)`,
132            /// so the keyset is stable across versions and callers can distinguish "unknown
133            /// key" from "present but unset".
134            pub fn render<F>(
135                &self,
136                meter: &mut impl ::mysten_common::rpc_format::Meter,
137            ) -> ::std::result::Result<
138                std::collections::BTreeMap<String, F>,
139                ::mysten_common::rpc_format::MeterError,
140            >
141            where
142                F: ::mysten_common::rpc_format::Format,
143            {
144                let mut map = std::collections::BTreeMap::new();
145                #(#render_arms)*
146                Ok(map)
147            }
148
149            /// Get the feature flags
150            pub fn lookup_feature(&self, value: String) -> Option<bool> {
151                self.feature_flags.lookup_attr(value)
152            }
153
154            pub fn feature_map(&self) -> std::collections::BTreeMap<String, bool> {
155                self.feature_flags.attr_map()
156            }
157        }
158
159        impl #struct_name {
160            #(#setters)*
161
162            pub fn set_attr_for_testing(&mut self, attr: String, val: String) {
163                match attr.as_str() {
164                    #(#scalar_value_setter_arms)*
165                    _ => panic!(
166                        "Attempting to set unknown or non-string-settable attribute: {}",
167                        attr,
168                    ),
169                }
170            }
171        }
172
173        #[allow(non_camel_case_types)]
174        #[derive(Clone, Serialize, Debug, PartialEq, Deserialize, schemars::JsonSchema)]
175        pub enum ProtocolConfigValue {
176            #(#variant_decls,)*
177        }
178
179        impl std::fmt::Display for ProtocolConfigValue {
180            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
181                use std::fmt::Write;
182                let mut writer = String::new();
183                match self {
184                    #(
185                        ProtocolConfigValue::#display_variants(x) => {
186                            write!(writer, "{}", x)?;
187                        }
188                    )*
189                }
190                write!(f, "{}", writer)
191            }
192        }
193    };
194
195    TokenStream::from(output)
196}
197
198/// Token streams emitted for a single `ProtocolConfig` field. Every field contributes an
199/// accessor, a setter, and a `render` arm; scalar fields additionally populate
200/// [`ScalarExtras`] for the `ProtocolConfigValue` / `attr_map` / `set_attr_for_testing` paths.
201struct ExpandedField {
202    /// `fn field_name(&self) -> T` (scalars only) + `fn field_name_as_option(&self) -> Option<T>`.
203    accessor: proc_macro2::TokenStream,
204    /// `set_x_for_testing` + `disable_x_for_testing` (always) and the `from_str` variant
205    /// (scalars only — non-scalars don't generally implement `FromStr`).
206    setter: proc_macro2::TokenStream,
207    /// One per-field block of `render` that calls `ToFormat::to_format` and inserts the result.
208    render_arm: proc_macro2::TokenStream,
209    /// Populated only when the field is a scalar (i.e. inner type is a single bare identifier
210    /// usable as a `ProtocolConfigValue` variant ident — `u16`/`u32`/`u64`/`bool`).
211    scalar_extras: Option<ScalarExtras>,
212}
213
214/// The extra tokens scalar fields contribute on top of the always-emitted accessor/setter/render
215/// pieces. Non-scalar fields don't participate in `ProtocolConfigValue` at all.
216struct ScalarExtras {
217    /// `stringify!(field_name) => self.set_x_from_str_for_testing(val),` — match arm for
218    /// `set_attr_for_testing`.
219    value_setter_arm: proc_macro2::TokenStream,
220    /// `stringify!(field_name) => self.field_name.map(ProtocolConfigValue::Variant),` — match
221    /// arm for `lookup_attr`.
222    lookup_arm: proc_macro2::TokenStream,
223    /// `stringify!(field_name)` — used to assemble `attr_map`.
224    field_name_str: proc_macro2::TokenStream,
225    /// Variant identifier used in `ProtocolConfigValue`.
226    variant_ident: syn::Ident,
227    /// Inner `T` of `Option<T>` — pairs with `variant_ident` when declaring the enum variant.
228    inner_type: syn::Type,
229}
230
231fn expand_field(f: &AccessorField) -> ExpandedField {
232    let field_name = &f.field_name;
233    let field_type = &f.field_type;
234    let inner_type = &f.inner_type;
235    let as_option_name: proc_macro2::TokenStream =
236        format!("{field_name}_as_option").parse().unwrap();
237    let test_setter_name: proc_macro2::TokenStream =
238        format!("set_{field_name}_for_testing").parse().unwrap();
239    let test_un_setter_name: proc_macro2::TokenStream =
240        format!("disable_{field_name}_for_testing").parse().unwrap();
241
242    let render_arm = quote! {
243        {
244            let value = match self.#field_name.as_ref() {
245                Some(v) => <_ as ::mysten_common::rpc_format::ToFormat>::to_format::<F, _>(v, meter)?,
246                None => <F as ::mysten_common::rpc_format::Format>::null(meter)?,
247            };
248            map.insert(stringify!(#field_name).to_owned(), value);
249        }
250    };
251
252    // `_as_option` is always emitted. The plain getter and the string-based setter only make
253    // sense for scalars (the plain getter unwraps to a `Copy` primitive; the `from_str` setter
254    // requires `FromStr` on the inner type). Non-scalars provide custom getters next to the
255    // field definition when they want one (e.g. borrowed-slice ergonomics).
256    let as_option_emit = quote! {
257        pub fn #as_option_name(&self) -> #field_type {
258            self.#field_name.clone()
259        }
260    };
261    let typed_setter = if f.custom_setter {
262        quote! {}
263    } else {
264        quote! {
265            pub fn #test_setter_name(&mut self, val: #inner_type) {
266                self.#field_name = Some(val);
267            }
268        }
269    };
270    let common_setters = quote! {
271        #typed_setter
272
273        pub fn #test_un_setter_name(&mut self) {
274            self.#field_name = None;
275        }
276    };
277
278    match &f.scalar_variant {
279        Some(variant_ident) => {
280            let test_setter_from_str_name: proc_macro2::TokenStream =
281                format!("set_{field_name}_from_str_for_testing")
282                    .parse()
283                    .unwrap();
284            ExpandedField {
285                accessor: quote! {
286                    pub fn #field_name(&self) -> #inner_type {
287                        self.#field_name.expect(Self::CONSTANT_ERR_MSG)
288                    }
289
290                    pub fn #as_option_name(&self) -> #field_type {
291                        self.#field_name
292                    }
293                },
294                setter: quote! {
295                    #common_setters
296
297                    pub fn #test_setter_from_str_name(&mut self, val: String) {
298                        use std::str::FromStr;
299                        self.#test_setter_name(#inner_type::from_str(&val).unwrap());
300                    }
301                },
302                render_arm,
303                scalar_extras: Some(ScalarExtras {
304                    value_setter_arm: quote! {
305                        stringify!(#field_name) => self.#test_setter_from_str_name(val),
306                    },
307                    lookup_arm: quote! {
308                        stringify!(#field_name) => self
309                            .#field_name
310                            .map(ProtocolConfigValue::#variant_ident),
311                    },
312                    field_name_str: quote! { stringify!(#field_name) },
313                    variant_ident: variant_ident.clone(),
314                    inner_type: inner_type.clone(),
315                }),
316            }
317        }
318        None => ExpandedField {
319            accessor: as_option_emit,
320            setter: common_setters,
321            render_arm,
322            scalar_extras: None,
323        },
324    }
325}
326
327/// Per-field metadata extracted from a `ProtocolConfig` field while expanding the
328/// `ProtocolConfigAccessors` derive.
329struct AccessorField {
330    /// The `#field_name` identifier — used both for accessor method names and as the string key
331    /// in the generated maps.
332    field_name: syn::Ident,
333    /// The full `Option<T>` type as written in the struct.
334    field_type: syn::Type,
335    /// The inner `T` extracted from `Option<T>`.
336    inner_type: syn::Type,
337    /// `Some(ident)` when the inner type is a single bare identifier usable directly as a
338    /// `ProtocolConfigValue` variant ident (`u16`/`u32`/`u64`/`bool`). `None` for non-scalar
339    /// fields, which never appear in `ProtocolConfigValue`.
340    scalar_variant: Option<syn::Ident>,
341    /// `#[custom_setter]` — the struct hand-writes `set_x_for_testing` instead of the derive
342    /// emitting the plain field assignment.
343    custom_setter: bool,
344}
345
346fn parse_accessor_field(field: &syn::Field) -> Option<AccessorField> {
347    let field_name = field.ident.clone().expect("Field must be named");
348
349    let skip_accessor = field
350        .attrs
351        .iter()
352        .any(|attr| attr.path.is_ident("skip_accessor"));
353    if skip_accessor {
354        return None;
355    }
356
357    let field_type = &field.ty;
358    let type_path = match field_type {
359        Type::Path(p) => p,
360        _ => return None,
361    };
362    let last_segment = type_path.path.segments.last()?;
363    if last_segment.ident != "Option" {
364        return None;
365    }
366    let inner_type = match &last_segment.arguments {
367        syn::PathArguments::AngleBracketed(args) => match args.args.first()? {
368            syn::GenericArgument::Type(ty) => ty.clone(),
369            _ => panic!("Expected a type argument inside Option<...> for `{field_name}`"),
370        },
371        _ => panic!("Expected angle bracketed arguments inside Option<...> for `{field_name}`"),
372    };
373
374    let scalar_variant = inferred_scalar_variant_ident(&inner_type);
375
376    let custom_setter = field
377        .attrs
378        .iter()
379        .any(|attr| attr.path.is_ident("custom_setter"));
380
381    Some(AccessorField {
382        field_name,
383        field_type: field_type.clone(),
384        inner_type,
385        scalar_variant,
386        custom_setter,
387    })
388}
389
390fn inferred_scalar_variant_ident(ty: &syn::Type) -> Option<syn::Ident> {
391    const SCALAR_PRIMITIVES: &[&str] = &["bool", "u8", "u16", "u32", "u64", "u128", "usize"];
392
393    let Type::Path(path) = ty else { return None };
394    if path.qself.is_some() {
395        return None;
396    }
397    if path.path.segments.len() != 1 {
398        return None;
399    }
400    let segment = path.path.segments.first()?;
401    if !matches!(segment.arguments, syn::PathArguments::None) {
402        return None;
403    }
404    if !SCALAR_PRIMITIVES.iter().any(|p| segment.ident == *p) {
405        return None;
406    }
407    Some(segment.ident.clone())
408}
409
410#[proc_macro_derive(ProtocolConfigOverride)]
411pub fn protocol_config_override_macro(input: TokenStream) -> TokenStream {
412    let ast = parse_macro_input!(input as DeriveInput);
413
414    // Create a new struct name by appending "Optional".
415    let struct_name = &ast.ident;
416    let optional_struct_name =
417        syn::Ident::new(&format!("{}Optional", struct_name), struct_name.span());
418
419    // Extract the fields from the struct
420    let fields = match &ast.data {
421        Data::Struct(data_struct) => match &data_struct.fields {
422            Fields::Named(fields_named) => &fields_named.named,
423            _ => panic!("ProtocolConfig must have named fields"),
424        },
425        _ => panic!("ProtocolConfig must be a struct"),
426    };
427
428    // Create new fields with types wrapped in Option.
429    let optional_fields = fields.iter().map(|field| {
430        let field_name = &field.ident;
431        let field_type = &field.ty;
432        quote! {
433            #field_name: Option<#field_type>
434        }
435    });
436
437    // Generate the function to update the original struct.
438    let update_fields = fields.iter().map(|field| {
439        let field_name = &field.ident;
440        quote! {
441            if let Some(value) = self.#field_name {
442                tracing::warn!(
443                    "ProtocolConfig field \"{}\" has been overridden with the value: {value:?}",
444                    stringify!(#field_name),
445                );
446                config.#field_name = value;
447            }
448        }
449    });
450
451    // Generate the new struct definition.
452    let output = quote! {
453        #[derive(serde::Deserialize, Debug)]
454        pub struct #optional_struct_name {
455            #(#optional_fields,)*
456        }
457
458        impl #optional_struct_name {
459            pub fn apply_to(self, config: &mut #struct_name) {
460                #(#update_fields)*
461            }
462        }
463    };
464
465    TokenStream::from(output)
466}
467
468#[proc_macro_derive(
469    ProtocolConfigFeatureFlagsGetters,
470    attributes(skip_accessor, skip_protocol_config_accessor)
471)]
472pub fn feature_flag_getters_macro(input: TokenStream) -> TokenStream {
473    let ast = parse_macro_input!(input as DeriveInput);
474
475    let struct_name = &ast.ident;
476    let data = &ast.data;
477
478    let getters = match data {
479        Data::Struct(data_struct) => match &data_struct.fields {
480            // Operate on each field of the ProtocolConfig struct
481            Fields::Named(fields_named) => fields_named
482                .named
483                .iter()
484                .filter_map(|field| {
485                    // Extract field name and type
486                    let field_name = field.ident.as_ref().expect("Field must be named");
487                    let field_type = &field.ty;
488                    let skip_accessor = field
489                        .attrs
490                        .iter()
491                        .any(|attr| attr.path.is_ident("skip_accessor"));
492                    if skip_accessor {
493                        return None;
494                    }
495                    let skip_protocol_config_accessor = field
496                        .attrs
497                        .iter()
498                        .any(|attr| attr.path.is_ident("skip_protocol_config_accessor"));
499                    // Check if field is of type bool
500                    match field_type {
501                        Type::Path(type_path)
502                            if type_path
503                                .path
504                                .segments
505                                .last()
506                                .is_some_and(|segment| segment.ident == "bool") =>
507                        {
508                            let getter = if skip_protocol_config_accessor {
509                                quote! {}
510                            } else {
511                                quote! {
512                                    // Forward the flag from ProtocolConfig.
513                                    pub fn #field_name(&self) -> #field_type {
514                                        self.feature_flags.#field_name
515                                    }
516                                }
517                            };
518                            let setter_name: proc_macro2::TokenStream =
519                                format!("set_{}_for_testing", field_name).parse().unwrap();
520                            let protocol_config_getter = quote! {
521                                #getter
522
523                                pub fn #setter_name(&mut self, val: bool) {
524                                    self.feature_flags.#field_name = val;
525                                }
526                            };
527                            Some((
528                                protocol_config_getter,
529                                (
530                                    quote! {
531                                        stringify!(#field_name) => Some(self.#field_name),
532                                    },
533                                    (
534                                        quote! {
535                                            stringify!(#field_name) => self.#field_name = val,
536                                        },
537                                        quote! {
538                                            stringify!(#field_name)
539                                        },
540                                    ),
541                                ),
542                            ))
543                        }
544                        _ => None,
545                    }
546                })
547                .collect::<Vec<_>>(),
548            _ => panic!("Only named fields are supported."),
549        },
550        _ => panic!("Only structs supported."),
551    };
552
553    let mut protocol_config_getters = Vec::new();
554    let mut string_name_getters = Vec::new();
555    let mut string_name_setters = Vec::new();
556    let mut field_names = Vec::new();
557    for (protocol_config_getter, (string_name_getter, (string_name_setter, field_name))) in getters
558    {
559        protocol_config_getters.push(protocol_config_getter);
560        string_name_getters.push(string_name_getter);
561        string_name_setters.push(string_name_setter);
562        field_names.push(field_name);
563    }
564
565    let output = quote! {
566        impl #struct_name {
567            /// Lookup a feature flag by its string representation
568            pub fn lookup_attr(&self, value: String) -> Option<bool> {
569                match value.as_str() {
570                    #(#string_name_getters)*
571                    _ => None,
572                }
573            }
574
575            /// Set a feature flag by its string representation
576            pub fn set_attr_for_testing(&mut self, attr: String, val: bool) {
577                match attr.as_str() {
578                    #(#string_name_setters)*
579                    _ => panic!("Attempting to set unknown feature flag: {}", attr),
580                }
581            }
582
583            /// Get a map of all feature flags from string representations
584            pub fn attr_map(&self) -> std::collections::BTreeMap<String, bool> {
585                vec![
586                    // Okay to unwrap since we added all above
587                    #(((#field_names).to_owned(), self.lookup_attr((#field_names).to_owned()).unwrap()),)*
588                    ].into_iter().collect()
589            }
590        }
591
592        impl ProtocolConfig {
593            #(#protocol_config_getters)*
594
595            /// Set a feature flag by its string representation
596            pub fn set_feature_flag_for_testing(&mut self, flag: String, val: bool) {
597                self.feature_flags.set_attr_for_testing(flag, val)
598            }
599        }
600    };
601
602    TokenStream::from(output)
603}
sui_protocol_config_macros/lib.rs

sui_protocol_config_macros/
lib.rs
