query_group_macro/
queries.rs

1//! The IR of the `#[query_group]` macro.
2
3use quote::{ToTokens, format_ident, quote, quote_spanned};
4use syn::{FnArg, Ident, PatType, Path, Receiver, ReturnType, Type, parse_quote, spanned::Spanned};
5
6use crate::Cycle;
7
8pub(crate) struct TrackedQuery {
9    pub(crate) trait_name: Ident,
10    pub(crate) signature: syn::Signature,
11    pub(crate) pat_and_tys: Vec<PatType>,
12    pub(crate) invoke: Option<Path>,
13    pub(crate) default: Option<syn::Block>,
14    pub(crate) cycle: Option<Cycle>,
15    pub(crate) lru: Option<u32>,
16    pub(crate) generated_struct: Option<GeneratedInputStruct>,
17}
18
19pub(crate) struct GeneratedInputStruct {
20    pub(crate) input_struct_name: Ident,
21    pub(crate) create_data_ident: Ident,
22}
23
24impl ToTokens for TrackedQuery {
25    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
26        let sig = &self.signature;
27        let trait_name = &self.trait_name;
28
29        let ret = &sig.output;
30
31        let invoke = match &self.invoke {
32            Some(path) => path.to_token_stream(),
33            None => sig.ident.to_token_stream(),
34        };
35
36        let fn_ident = &sig.ident;
37        let shim: Ident = format_ident!("{}_shim", fn_ident);
38
39        let options = self
40            .cycle
41            .as_ref()
42            .map(|Cycle { cycle_fn, cycle_initial, cycle_result }| {
43                let cycle_fn = cycle_fn.as_ref().map(|(ident, path)| quote!(#ident=#path));
44                let cycle_initial =
45                    cycle_initial.as_ref().map(|(ident, path)| quote!(#ident=#path));
46                let cycle_result = cycle_result.as_ref().map(|(ident, path)| quote!(#ident=#path));
47                let options = cycle_fn.into_iter().chain(cycle_initial).chain(cycle_result);
48                quote!(#(#options),*)
49            })
50            .into_iter()
51            .chain(self.lru.map(|lru| quote!(lru = #lru)))
52            .chain(Some(quote!(unsafe(non_update_return_type))));
53        let annotation = quote!(#[salsa_macros::tracked( #(#options),* )]);
54
55        let pat_and_tys = &self.pat_and_tys;
56        let params = self
57            .pat_and_tys
58            .iter()
59            .map(|pat_type| pat_type.pat.clone())
60            .collect::<Vec<Box<syn::Pat>>>();
61
62        let invoke_block = match &self.default {
63            Some(default) => quote! { #default },
64            None => {
65                let invoke_params: proc_macro2::TokenStream = quote! {db, #(#params),*};
66                quote_spanned! { invoke.span() =>  {#invoke(#invoke_params)}}
67            }
68        };
69
70        let method = match &self.generated_struct {
71            Some(generated_struct) => {
72                let input_struct_name = &generated_struct.input_struct_name;
73                let create_data_ident = &generated_struct.create_data_ident;
74
75                quote! {
76                    #sig {
77                        #annotation
78                        fn #shim<'db>(
79                            db: &'db dyn #trait_name,
80                            _input: #input_struct_name,
81                            #(#pat_and_tys),*
82                        ) #ret
83                            #invoke_block
84                        #shim(self, #create_data_ident(self), #(#params),*)
85                    }
86                }
87            }
88            None => {
89                quote! {
90                    #sig {
91                        #annotation
92                        fn #shim<'db>(
93                            db: &'db dyn #trait_name,
94                            #(#pat_and_tys),*
95                        ) #ret
96                            #invoke_block
97
98                        #shim(self, #(#params),*)
99                    }
100                }
101            }
102        };
103
104        method.to_tokens(tokens);
105    }
106}
107
108pub(crate) struct InputQuery {
109    pub(crate) signature: syn::Signature,
110    pub(crate) create_data_ident: Ident,
111}
112
113impl ToTokens for InputQuery {
114    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
115        let sig = &self.signature;
116        let fn_ident = &sig.ident;
117        let create_data_ident = &self.create_data_ident;
118
119        let method = quote! {
120            #sig {
121                let data = #create_data_ident(self);
122                data.#fn_ident(self).unwrap()
123            }
124        };
125        method.to_tokens(tokens);
126    }
127}
128
129pub(crate) struct InputSetter {
130    pub(crate) signature: syn::Signature,
131    pub(crate) return_type: syn::Type,
132    pub(crate) create_data_ident: Ident,
133}
134
135impl ToTokens for InputSetter {
136    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
137        let sig = &mut self.signature.clone();
138
139        let ty = &self.return_type;
140        let fn_ident = &sig.ident;
141        let create_data_ident = &self.create_data_ident;
142
143        let setter_ident = format_ident!("set_{}", fn_ident);
144        sig.ident = setter_ident.clone();
145
146        let value_argument: PatType = parse_quote!(__value: #ty);
147        sig.inputs.push(FnArg::Typed(value_argument.clone()));
148
149        // make `&self` `&mut self` instead.
150        let mut_receiver: Receiver = parse_quote!(&mut self);
151        if let Some(og) = sig.inputs.first_mut() {
152            *og = FnArg::Receiver(mut_receiver)
153        }
154
155        // remove the return value.
156        sig.output = ReturnType::Default;
157
158        let value = &value_argument.pat;
159        let method = quote! {
160            #sig {
161                use salsa::Setter;
162                let data = #create_data_ident(self);
163                data.#setter_ident(self).to(Some(#value));
164            }
165        };
166        method.to_tokens(tokens);
167    }
168}
169
170pub(crate) struct InputSetterWithDurability {
171    pub(crate) signature: syn::Signature,
172    pub(crate) return_type: syn::Type,
173    pub(crate) create_data_ident: Ident,
174}
175
176impl ToTokens for InputSetterWithDurability {
177    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
178        let sig = &mut self.signature.clone();
179
180        let ty = &self.return_type;
181        let fn_ident = &sig.ident;
182        let setter_ident = format_ident!("set_{}", fn_ident);
183
184        let create_data_ident = &self.create_data_ident;
185
186        sig.ident = format_ident!("set_{}_with_durability", fn_ident);
187
188        let value_argument: PatType = parse_quote!(__value: #ty);
189        sig.inputs.push(FnArg::Typed(value_argument.clone()));
190
191        let durability_argument: PatType = parse_quote!(durability: salsa::Durability);
192        sig.inputs.push(FnArg::Typed(durability_argument.clone()));
193
194        // make `&self` `&mut self` instead.
195        let mut_receiver: Receiver = parse_quote!(&mut self);
196        if let Some(og) = sig.inputs.first_mut() {
197            *og = FnArg::Receiver(mut_receiver)
198        }
199
200        // remove the return value.
201        sig.output = ReturnType::Default;
202
203        let value = &value_argument.pat;
204        let durability = &durability_argument.pat;
205        let method = quote! {
206            #sig {
207                use salsa::Setter;
208                let data = #create_data_ident(self);
209                data.#setter_ident(self)
210                    .with_durability(#durability)
211                    .to(Some(#value));
212            }
213        };
214        method.to_tokens(tokens);
215    }
216}
217
218pub(crate) enum SetterKind {
219    Plain(InputSetter),
220    WithDurability(InputSetterWithDurability),
221}
222
223impl ToTokens for SetterKind {
224    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
225        match self {
226            SetterKind::Plain(input_setter) => input_setter.to_tokens(tokens),
227            SetterKind::WithDurability(input_setter_with_durability) => {
228                input_setter_with_durability.to_tokens(tokens)
229            }
230        }
231    }
232}
233
234pub(crate) struct Transparent {
235    pub(crate) signature: syn::Signature,
236    pub(crate) pat_and_tys: Vec<PatType>,
237    pub(crate) invoke: Option<Path>,
238    pub(crate) default: Option<syn::Block>,
239}
240
241impl ToTokens for Transparent {
242    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
243        let sig = &self.signature;
244
245        let ty = self
246            .pat_and_tys
247            .iter()
248            .map(|pat_type| pat_type.pat.clone())
249            .collect::<Vec<Box<syn::Pat>>>();
250
251        let invoke = match &self.invoke {
252            Some(path) => path.to_token_stream(),
253            None => sig.ident.to_token_stream(),
254        };
255
256        let method = match &self.default {
257            Some(default) => quote! {
258                #sig { let db = self; #default }
259            },
260            None => quote! {
261                #sig {
262                    #invoke(self, #(#ty),*)
263                }
264            },
265        };
266
267        method.to_tokens(tokens);
268    }
269}
270pub(crate) struct Intern {
271    pub(crate) signature: syn::Signature,
272    pub(crate) pat_and_tys: Vec<PatType>,
273    pub(crate) interned_struct_path: Path,
274}
275
276impl ToTokens for Intern {
277    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
278        let sig = &self.signature;
279
280        let ty = self.pat_and_tys.to_vec();
281
282        let interned_pat = ty.first().expect("at least one pat; this is a bug");
283        let interned_pat = &interned_pat.pat;
284
285        let wrapper_struct = self.interned_struct_path.to_token_stream();
286
287        let method = quote! {
288            #sig {
289                #wrapper_struct::new(self, #interned_pat)
290            }
291        };
292
293        method.to_tokens(tokens);
294    }
295}
296
297pub(crate) struct Lookup {
298    pub(crate) signature: syn::Signature,
299    pub(crate) pat_and_tys: Vec<PatType>,
300    pub(crate) return_ty: Type,
301    pub(crate) interned_struct_path: Path,
302}
303
304impl Lookup {
305    pub(crate) fn prepare_signature(&mut self) {
306        let sig = &self.signature;
307
308        let ident = format_ident!("lookup_{}", sig.ident);
309
310        let ty = self.pat_and_tys.to_vec();
311
312        let interned_key = &self.return_ty;
313
314        let interned_pat = ty.first().expect("at least one pat; this is a bug");
315        let interned_return_ty = &interned_pat.ty;
316
317        self.signature = parse_quote!(
318            fn #ident(&self, id: #interned_key) -> #interned_return_ty
319        );
320    }
321}
322
323impl ToTokens for Lookup {
324    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
325        let sig = &self.signature;
326
327        let wrapper_struct = self.interned_struct_path.to_token_stream();
328        let method = quote! {
329            #sig {
330                let zalsa = self.zalsa();
331                #wrapper_struct::ingredient(zalsa).data(zalsa, id.as_id()).0.clone()
332            }
333        };
334
335        method.to_tokens(tokens);
336    }
337}
338
339#[allow(clippy::large_enum_variant)]
340pub(crate) enum Queries {
341    TrackedQuery(TrackedQuery),
342    InputQuery(InputQuery),
343    Intern(Intern),
344    Transparent(Transparent),
345}
346
347impl ToTokens for Queries {
348    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
349        match self {
350            Queries::TrackedQuery(tracked_query) => tracked_query.to_tokens(tokens),
351            Queries::InputQuery(input_query) => input_query.to_tokens(tokens),
352            Queries::Transparent(transparent) => transparent.to_tokens(tokens),
353            Queries::Intern(intern) => intern.to_tokens(tokens),
354        }
355    }
356}
query_group_macro/queries.rs

query_group_macro/
queries.rs
