Skip to main content

vls_policy_derive/
lib.rs

1use proc_macro::TokenStream;
2use quote::{format_ident, quote};
3use syn::{parse_macro_input, Data, DeriveInput, Fields};
4
5/// Derives a new struct and a merge function for a given struct.
6/// The new struct has the same fields as the original, but each field is wrapped in an `Option`.
7/// See generated `merge` and `resolve_defaults` functions for more information.
8/// The struct must implement `Clone`.
9#[proc_macro_derive(Optionized)]
10pub fn derive_optionized(input: TokenStream) -> TokenStream {
11    let ast = parse_macro_input!(input as DeriveInput);
12    let struct_name = &ast.ident;
13
14    let fields: Vec<_> = match &ast.data {
15        Data::Struct(data) => match &data.fields {
16            Fields::Named(fields) => fields
17                .named
18                .iter()
19                .map(|field| {
20                    let field_name = field.ident.clone().unwrap();
21                    let field_type = &field.ty;
22                    (field_name, field_type)
23                })
24                .collect(),
25            _ => return quote! {
26                compile_error!("SimplePolicyMerge can only be derived for structs with named fields");
27            }.into(),
28        },
29        _ => return quote! {
30            compile_error!("SimplePolicyMerge can only be derived for structs");
31        }.into(),
32    };
33
34    let optionized_struct_name = format_ident!("Optionized{}", struct_name);
35    let optionized_fields = fields.iter().map(|(field_name, field_type)| {
36        quote! {
37            pub #field_name: ::core::option::Option<#field_type>
38        }
39    });
40    let optionized_struct = quote! {
41        /// An optionized version of #struct_name.
42        /// Each field is wrapped in an Option.
43        /// See `merge` and `resolve_defaults` for more information.
44        #[derive(Deserialize, Debug, Default, Clone)]
45        #[allow(missing_docs)]
46        pub struct #optionized_struct_name {
47            #(#optionized_fields),*
48        }
49    };
50
51    let merge_impl = generate_merge_function(&optionized_struct_name, &fields);
52    let resolve_defaults_function =
53        generate_resolve_defaults_function(&optionized_struct_name, &fields);
54    let new_function = generate_new_function(&optionized_struct_name, &fields);
55
56    let expanded = quote! {
57        #optionized_struct
58        #new_function
59        #merge_impl
60        #resolve_defaults_function
61    };
62
63    expanded.into()
64}
65
66/// Generates a merge function for the given struct.
67/// The merge function takes another instance of the struct as input,
68/// and for each field, if the input has a `Some` value for that field, it sets the field in `self` to that value.
69fn generate_merge_function(
70    struct_name: &proc_macro2::Ident,
71    fields: &[(proc_macro2::Ident, &syn::Type)],
72) -> proc_macro2::TokenStream {
73    let field_assignments = fields.iter().map(|(field_name, _)| {
74        quote! {
75            if let Some(val) = other.#field_name {
76                self.#field_name = Some(val);
77            }
78        }
79    });
80    quote! {
81        impl #struct_name {
82            /// Merge all fields from `other` into `self`: for each field,
83            /// if `other` has a `Some` value for that field, set the field in `self` to that value.
84            pub fn merge(&mut self, other: Self) {
85                #(#field_assignments)*
86            }
87        }
88    }
89}
90
91fn generate_resolve_defaults_function(
92    struct_name: &proc_macro2::Ident,
93    fields: &[(proc_macro2::Ident, &syn::Type)],
94) -> proc_macro2::TokenStream {
95    let binding = struct_name.to_string();
96    let original_struct_name = binding.trim_start_matches("Optionized");
97    let original_struct_ident = format_ident!("{}", original_struct_name);
98
99    let field_assignments = fields.iter().map(|(field_name, _)| {
100        quote! {
101            #field_name: self.#field_name.clone().unwrap_or(defaults.#field_name),
102        }
103    });
104    quote! {
105        impl #struct_name {
106            /// Apply defaults to fields in self that are None and return a new instance of the original struct
107            pub fn resolve_defaults(self, defaults: #original_struct_ident) -> #original_struct_ident {
108                #original_struct_ident {
109                    #(#field_assignments)*
110                }
111            }
112        }
113    }
114}
115
116fn generate_new_function(
117    struct_name: &proc_macro2::Ident,
118    fields: &[(proc_macro2::Ident, &syn::Type)],
119) -> proc_macro2::TokenStream {
120    let field_initializations = fields.iter().map(|(field_name, _)| {
121        quote! {
122            #field_name: None
123        }
124    });
125    quote! {
126        impl #struct_name {
127            /// Create a new #struct_name with all fields set to None
128            pub fn new() -> Self {
129                Self {
130                    #(#field_initializations),*
131                }
132            }
133        }
134    }
135}