Skip to main content

gproxy_protocol_macros/
lib.rs

1//! Derives and expression helpers used by `gproxy-protocol`.
2
3use proc_macro::TokenStream;
4use quote::{format_ident, quote};
5use syn::visit_mut::{self, VisitMut};
6use syn::{Data, DeriveInput, Expr, ExprStruct, Fields, Meta, Type, parse_macro_input};
7
8#[proc_macro_derive(WireBuilder, attributes(serde))]
9pub fn derive_wire_builder(input: TokenStream) -> TokenStream {
10    let input = parse_macro_input!(input as DeriveInput);
11    let name = input.ident;
12    let builder = format_ident!("Wire{name}Builder");
13    let generics = input.generics;
14    let (impl_generics, type_generics, where_clause) = generics.split_for_impl();
15    let Data::Struct(data) = input.data else {
16        return syn::Error::new_spanned(name, "WireBuilder only supports structs")
17            .to_compile_error()
18            .into();
19    };
20    let Fields::Named(fields) = data.fields else {
21        return syn::Error::new_spanned(name, "WireBuilder requires named fields")
22            .to_compile_error()
23            .into();
24    };
25
26    let fields: Vec<_> = fields.named.into_iter().collect();
27    let names: Vec<_> = fields
28        .iter()
29        .map(|field| field.ident.as_ref().expect("named field"))
30        .collect();
31    let types: Vec<_> = fields.iter().map(|field| &field.ty).collect();
32    let values = fields.iter().map(|field| {
33        let field_name = field.ident.as_ref().expect("named field");
34        if field_has_default(field) {
35            quote! { self.#field_name.unwrap_or_default() }
36        } else {
37            quote! {
38                self.#field_name.ok_or_else(|| crate::WireBuildError::missing(
39                    stringify!(#name),
40                    stringify!(#field_name),
41                ))?
42            }
43        }
44    });
45
46    quote! {
47        #[doc(hidden)]
48        pub struct #builder #generics {
49            #(#names: ::core::option::Option<#types>,)*
50        }
51
52        impl #impl_generics #builder #type_generics #where_clause {
53            #(
54                pub fn #names(mut self, value: #types) -> Self {
55                    self.#names = ::core::option::Option::Some(value);
56                    self
57                }
58            )*
59
60            pub fn build(self) -> ::core::result::Result<#name #type_generics, crate::WireBuildError> {
61                ::core::result::Result::Ok(#name {
62                    #(#names: #values,)*
63                })
64            }
65        }
66
67        impl #impl_generics #name #type_generics #where_clause {
68            pub fn builder() -> #builder #type_generics {
69                #builder {
70                    #(#names: ::core::option::Option::None,)*
71                }
72            }
73        }
74    }
75    .into()
76}
77
78fn field_has_default(field: &syn::Field) -> bool {
79    if matches!(
80        &field.ty,
81        Type::Path(path) if path.path.segments.last().is_some_and(|segment| segment.ident == "Option")
82    ) {
83        return true;
84    }
85    field.attrs.iter().any(|attribute| {
86        if !attribute.path().is_ident("serde") {
87            return false;
88        }
89        let Meta::List(list) = &attribute.meta else {
90            return false;
91        };
92        let tokens = list.tokens.to_string();
93        tokens.split(',').any(|part| {
94            let part = part.trim();
95            part == "default"
96                || part.starts_with("default =")
97                || part == "skip"
98                || part == "skip_deserializing"
99        })
100    })
101}
102
103/// Construct a non-exhaustive wire struct through its generated builder while
104/// retaining familiar named-field syntax at transform implementation sites.
105#[proc_macro]
106pub fn wire(input: TokenStream) -> TokenStream {
107    let expression = parse_macro_input!(input as ExprStruct);
108    expand_struct(expression).into()
109}
110
111struct NestedWireRewriter;
112
113impl VisitMut for NestedWireRewriter {
114    fn visit_expr_mut(&mut self, expression: &mut Expr) {
115        visit_mut::visit_expr_mut(self, expression);
116        let Expr::Struct(struct_expression) = expression else {
117            return;
118        };
119        if !looks_like_enum_variant(&struct_expression.path) {
120            *expression = Expr::Verbatim(expand_struct(struct_expression.clone()));
121        }
122    }
123}
124
125fn looks_like_enum_variant(path: &syn::Path) -> bool {
126    path.segments
127        .iter()
128        .rev()
129        .nth(1)
130        .and_then(|segment| segment.ident.to_string().chars().next())
131        .is_some_and(char::is_uppercase)
132}
133
134fn expand_struct(mut expression: ExprStruct) -> proc_macro2::TokenStream {
135    let mut rewriter = NestedWireRewriter;
136    for field in &mut expression.fields {
137        rewriter.visit_expr_mut(&mut field.expr);
138    }
139    if let Some(rest) = expression.rest.as_mut() {
140        rewriter.visit_expr_mut(rest);
141    }
142    if looks_like_enum_variant(&expression.path) {
143        return quote! { #expression };
144    }
145    let path = expression.path;
146    let fields = expression.fields;
147
148    if let Some(rest) = expression.rest {
149        let assignments = fields.iter().map(|field| {
150            let member = &field.member;
151            let value = &field.expr;
152            quote! { value.#member = #value; }
153        });
154        return quote! {{
155            let mut value: #path = #rest;
156            #(#assignments)*
157            value
158        }};
159    }
160
161    let setters = fields.iter().map(|field| {
162        let member = &field.member;
163        let value = &field.expr;
164        quote! { .#member(#value) }
165    });
166    quote! {{
167        #path::builder()
168            #(#setters)*
169            .build()
170            .expect(concat!("complete ", stringify!(#path), " wire construction"))
171    }}
172}