polyglot_sql_ast_derive/
lib.rs

1//! Internal derive used to keep Polyglot's typed AST traversal exhaustive.
2//!
3//! Fields are visited in declaration order. Standard containers are traversed
4//! automatically, while `#[ast(skip)]` excludes derived semantic metadata that
5//! is not part of the syntax tree.
6
7use proc_macro::TokenStream;
8use quote::{format_ident, quote};
9use syn::{
10    parse_macro_input, Data, DeriveInput, Field, Fields, GenericArgument, PathArguments, Type,
11};
12
13#[proc_macro_derive(AstNode, attributes(ast))]
14/// Derive immutable and mutable expression-child visitors for an AST payload.
15pub fn derive_ast_node(input: TokenStream) -> TokenStream {
16    let input = parse_macro_input!(input as DeriveInput);
17    expand_ast_node(&input).into()
18}
19
20fn expand_ast_node(input: &DeriveInput) -> proc_macro2::TokenStream {
21    let name = &input.ident;
22    let immutable = match &input.data {
23        Data::Struct(data) => visit_fields(&data.fields, false),
24        Data::Enum(data) => {
25            let arms = data.variants.iter().map(|variant| {
26                let variant_name = &variant.ident;
27                let (pattern, body) =
28                    visit_variant_fields(name, variant_name, &variant.fields, false);
29                quote!(#pattern => { #body })
30            });
31            quote!(match self { #(#arms),* })
32        }
33        Data::Union(_) => quote!(),
34    };
35    let mutable = match &input.data {
36        Data::Struct(data) => visit_fields(&data.fields, true),
37        Data::Enum(data) => {
38            let arms = data.variants.iter().map(|variant| {
39                let variant_name = &variant.ident;
40                let (pattern, body) =
41                    visit_variant_fields(name, variant_name, &variant.fields, true);
42                quote!(#pattern => { #body })
43            });
44            quote!(match self { #(#arms),* })
45        }
46        Data::Union(_) => quote!(),
47    };
48
49    quote! {
50        impl crate::ast_children::AstNode for #name {
51            fn visit_expressions<'ast, F>(
52                &'ast self,
53                path: &mut Vec<crate::ast_children::ChildPathSegment>,
54                visitor: &mut F,
55            )
56            where
57                F: FnMut(
58                    &[crate::ast_children::ChildPathSegment],
59                    &'ast crate::expressions::Expression,
60                ),
61            {
62                #immutable
63            }
64
65            fn visit_expressions_mut<F>(
66                &mut self,
67                visitor: &mut F,
68            )
69            where
70                F: FnMut(&mut crate::expressions::Expression),
71            {
72                #mutable
73            }
74        }
75    }
76}
77
78fn visit_fields(fields: &Fields, mutable: bool) -> proc_macro2::TokenStream {
79    match fields {
80        Fields::Named(fields) => {
81            let visits = fields.named.iter().filter_map(|field| {
82                if is_skipped(field) {
83                    return None;
84                }
85                let ident = field.ident.as_ref().expect("named field");
86                let field_name = ident.to_string();
87                let access = quote!(&mut self.#ident);
88                let immutable_access = quote!(&self.#ident);
89                Some(if mutable {
90                    mutable_visit(&field.ty, access)
91                } else {
92                    let visit = immutable_visit(&field.ty, immutable_access);
93                    quote! {
94                        path.push(crate::ast_children::ChildPathSegment::Field(#field_name));
95                        #visit
96                        path.pop();
97                    }
98                })
99            });
100            quote!(#(#visits)*)
101        }
102        Fields::Unnamed(fields) => {
103            let visits = fields
104                .unnamed
105                .iter()
106                .enumerate()
107                .filter_map(|(index, field)| {
108                    if is_skipped(field) {
109                        return None;
110                    }
111                    let index = syn::Index::from(index);
112                    let access = quote!(&mut self.#index);
113                    let immutable_access = quote!(&self.#index);
114                    Some(if mutable {
115                        mutable_visit(&field.ty, access)
116                    } else {
117                        let visit = immutable_visit(&field.ty, immutable_access);
118                        quote! {
119                            path.push(crate::ast_children::ChildPathSegment::Index(#index));
120                            #visit
121                            path.pop();
122                        }
123                    })
124                });
125            quote!(#(#visits)*)
126        }
127        Fields::Unit => quote!(),
128    }
129}
130
131fn visit_variant_fields(
132    enum_name: &syn::Ident,
133    variant_name: &syn::Ident,
134    fields: &Fields,
135    mutable: bool,
136) -> (proc_macro2::TokenStream, proc_macro2::TokenStream) {
137    match fields {
138        Fields::Named(fields) => {
139            let bindings: Vec<_> = fields
140                .named
141                .iter()
142                .map(|field| {
143                    let ident = field.ident.as_ref().expect("named field");
144                    if is_skipped(field) || is_scalar(&field.ty) {
145                        quote!(#ident: _)
146                    } else {
147                        quote!(#ident)
148                    }
149                })
150                .collect();
151            let visits = fields.named.iter().filter_map(|field| {
152                if is_skipped(field) {
153                    return None;
154                }
155                let ident = field.ident.as_ref().expect("named field");
156                let field_name = ident.to_string();
157                Some(if mutable {
158                    mutable_visit(&field.ty, quote!(#ident))
159                } else {
160                    let visit = immutable_visit(&field.ty, quote!(#ident));
161                    quote! {
162                        path.push(crate::ast_children::ChildPathSegment::Field(#field_name));
163                        #visit
164                        path.pop();
165                    }
166                })
167            });
168            (
169                quote!(#enum_name::#variant_name { #(#bindings),* }),
170                quote!(#(#visits)*),
171            )
172        }
173        Fields::Unnamed(fields) => {
174            let bindings: Vec<_> = (0..fields.unnamed.len())
175                .map(|index| format_ident!("field_{index}"))
176                .collect();
177            let single_expression_payload = enum_name == "Expression" && fields.unnamed.len() == 1;
178            let visits = fields
179                .unnamed
180                .iter()
181                .zip(bindings.iter())
182                .enumerate()
183                .filter_map(|(index, (field, binding))| {
184                    if is_skipped(field) {
185                        return None;
186                    }
187                    Some(if mutable {
188                        mutable_visit(&field.ty, quote!(#binding))
189                    } else {
190                        let visit = immutable_visit(&field.ty, quote!(#binding));
191                        if single_expression_payload {
192                            visit
193                        } else {
194                            quote! {
195                                path.push(crate::ast_children::ChildPathSegment::Index(#index));
196                                #visit
197                                path.pop();
198                            }
199                        }
200                    })
201                });
202            (
203                quote!(#enum_name::#variant_name(#(#bindings),*)),
204                quote!(#(#visits)*),
205            )
206        }
207        Fields::Unit => (quote!(#enum_name::#variant_name), quote!()),
208    }
209}
210
211fn immutable_visit(ty: &Type, access: proc_macro2::TokenStream) -> proc_macro2::TokenStream {
212    if is_expression(ty) {
213        return quote!(visitor(path, #access););
214    }
215    if let Some(inner) = container_inner(ty, "Option") {
216        let visit = immutable_visit(inner, quote!(value));
217        return quote!(if let Some(value) = (#access).as_ref() { #visit });
218    }
219    if let Some(inner) = container_inner(ty, "Box") {
220        return immutable_visit(inner, quote!((#access).as_ref()));
221    }
222    if let Some(inner) = container_inner(ty, "Vec") {
223        let visit = immutable_visit(inner, quote!(value));
224        return quote! {
225            for (index, value) in (#access).iter().enumerate() {
226                path.push(crate::ast_children::ChildPathSegment::Index(index));
227                #visit
228                path.pop();
229            }
230        };
231    }
232    if let Type::Tuple(tuple) = ty {
233        let visits = tuple.elems.iter().enumerate().map(|(index, element)| {
234            let tuple_index = syn::Index::from(index);
235            let visit = immutable_visit(element, quote!(&(#access).#tuple_index));
236            quote! {
237                path.push(crate::ast_children::ChildPathSegment::Index(#index));
238                #visit
239                path.pop();
240            }
241        });
242        return quote!(#(#visits)*);
243    }
244    if is_scalar(ty) {
245        return quote!();
246    }
247    quote!(crate::ast_children::AstNode::visit_expressions(#access, path, visitor);)
248}
249
250fn mutable_visit(ty: &Type, access: proc_macro2::TokenStream) -> proc_macro2::TokenStream {
251    if is_expression(ty) {
252        return quote!(visitor(#access););
253    }
254    if let Some(inner) = container_inner(ty, "Option") {
255        let visit = mutable_visit(inner, quote!(value));
256        return quote!(if let Some(value) = (#access).as_mut() { #visit });
257    }
258    if let Some(inner) = container_inner(ty, "Box") {
259        return mutable_visit(inner, quote!((#access).as_mut()));
260    }
261    if let Some(inner) = container_inner(ty, "Vec") {
262        let visit = mutable_visit(inner, quote!(value));
263        return quote!(for value in (#access).iter_mut() { #visit });
264    }
265    if let Type::Tuple(tuple) = ty {
266        let visits = tuple.elems.iter().enumerate().map(|(index, element)| {
267            let tuple_index = syn::Index::from(index);
268            mutable_visit(element, quote!(&mut (#access).#tuple_index))
269        });
270        return quote!(#(#visits)*);
271    }
272    if is_scalar(ty) {
273        return quote!();
274    }
275    quote!(crate::ast_children::AstNode::visit_expressions_mut(#access, visitor);)
276}
277
278fn container_inner<'a>(ty: &'a Type, expected: &str) -> Option<&'a Type> {
279    let Type::Path(path) = ty else { return None };
280    let segment = path.path.segments.last()?;
281    if segment.ident != expected {
282        return None;
283    }
284    let PathArguments::AngleBracketed(arguments) = &segment.arguments else {
285        return None;
286    };
287    arguments.args.iter().find_map(|argument| match argument {
288        GenericArgument::Type(ty) => Some(ty),
289        _ => None,
290    })
291}
292
293fn is_expression(ty: &Type) -> bool {
294    matches!(ty, Type::Path(path) if path.path.segments.last().is_some_and(|segment| segment.ident == "Expression"))
295}
296
297fn is_scalar(ty: &Type) -> bool {
298    let Type::Path(path) = ty else { return false };
299    let Some(segment) = path.path.segments.last() else {
300        return false;
301    };
302    matches!(
303        segment.ident.to_string().as_str(),
304        "bool"
305            | "char"
306            | "str"
307            | "String"
308            | "usize"
309            | "isize"
310            | "u8"
311            | "u16"
312            | "u32"
313            | "u64"
314            | "u128"
315            | "i8"
316            | "i16"
317            | "i32"
318            | "i64"
319            | "i128"
320            | "f32"
321            | "f64"
322            | "Span"
323    )
324}
325
326fn is_skipped(field: &Field) -> bool {
327    field.attrs.iter().any(|attribute| {
328        attribute.path().is_ident("ast")
329            && attribute
330                .parse_args::<syn::Ident>()
331                .is_ok_and(|ident| ident == "skip")
332    })
333}
polyglot_sql_ast_derive/lib.rs

polyglot_sql_ast_derive/
lib.rs
